`core.ssm.gated_delta_net.common`#

Module Contents#

Classes#

`GatedDeltaNetSubmodules`	Contains the module specs for the input linear, output norm, and output linear layers.
`GatedDeltaRuleInterface`	Unified typing protocol for GDN core computation interfaces.
`_GDNBase`	Common base class for the Gated Delta Net (GDN) family of layers.

Functions#

`_build_thd_cp_a2a_perm`
`_build_head_perm_for_split_sections`
`get_parameter_local_cp`	Get the local parameter for the current context parallel rank.
`tensor_a2a_cp2hp`	All-to-all context parallel to hidden parallel.
`tensor_a2a_hp2cp`	All-to-all hidden parallel to context parallel.
`torch_chunk_gated_delta_rule`	Torch-native implementation of chunked gated delta rule for deterministic mode. Need this because FLA is not deterministic.

Data#

logger

API#

core.ssm.gated_delta_net.common.logger#: ‘getLogger(…)’

class core.ssm.gated_delta_net.common.GatedDeltaNetSubmodules#

Contains the module specs for the input linear, output norm, and output linear layers.

in_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

out_norm: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

out_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

class core.ssm.gated_delta_net.common.GatedDeltaRuleInterface#

Bases: typing.Protocol

Unified typing protocol for GDN core computation interfaces.

__call__(

q: torch.Tensor,

k: torch.Tensor,

v: torch.Tensor,

g: torch.Tensor,

scale: float | None = None,

initial_state: torch.Tensor | None = None,

output_final_state: bool = False,

use_qk_l2norm_in_kernel: bool = False,

cu_seqlens: torch.LongTensor | None = None,

**kwargs,

) → tuple[torch.Tensor, torch.Tensor | None]#

class core.ssm.gated_delta_net.common._GDNBase( config: megatron.core.transformer.TransformerConfig, submodules: core.ssm.gated_delta_net.common.GatedDeltaNetSubmodules, layer_number: int = None, bias: bool = False, conv_bias: bool = False, conv_init: float | None = None, use_qk_l2norm: bool = True, A_init_range: tuple[float, float] = (1, 16), pg_collection: megatron.core.process_groups_config.ProcessGroupCollection = None, *, name: str | None = None, cp_comm_type: str | None = None, )#

Bases: megatron.core.transformer.module.MegatronModule

Common base class for the Gated Delta Net (GDN) family of layers.

Hosts everything the GDN variants share: the fused input projection, causal convolution on q/k/v, the CP all-to-all plumbing, the kernel-input preparation skeleton, the gated output norm + projection, and sharded checkpointing.

Initialization

Parameters:

config – The config of the model.
submodules – Contains the module specs for the input and output linear layers.
layer_number – The layer number of this GDN layer.
bias – Whether to use bias in the linear layers.
conv_bias – Whether to use bias in the causal convolution.
conv_init – The initialization range for the causal convolution weights.
use_qk_l2norm – Whether to use L2 normalization in the kernel of the gated delta rule.
A_init_range – The initialization range for the attention weights.
pg_collection – The required process groups to use for tensor model parallel and context parallel.
name (str | None) – module instance name passed top-down from its paranet module
cp_comm_type (Optional[str]) – Accepted for TransformerLayer compatibility and ignored; GDN implements context parallelism with its own all-to-alls rather than the attention CP communication schemes.

dt_bias_dim: int#: None

a_log_dim: int#: None

in_proj_qkvg_dim: int#: None

in_proj_extra_dim: int#: None

in_proj_dim: int#: None

dt_bias: torch.nn.Parameter#: None

A_log: torch.nn.Parameter#: None

gated_delta_rule: core.ssm.gated_delta_net.common.GatedDeltaRuleInterface#: None

abstractmethod _setup_variant_attrs()#

Set variant specifics on the module. Called once from __init__.

Must set:

in_proj_dim
in_proj_split_names
in_proj_split_sections
feat_dim_split
dt_bias_dim / a_log_dim (sizes of the gate parameters, which the base class creates after the conv1d module to preserve the original parameter registration order)
gated_delta_rule (the kernel callable).

reset_parameters()#: Reset the parameters.

abstractmethod forward(

hidden_states: torch.Tensor,

attention_mask: torch.Tensor,

inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None,

packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,

sequence_len_offset: Optional[int] = None,

*,

inference_params: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None,

**kwargs,

) → tuple[torch.Tensor, torch.Tensor]#

_gated_norm_and_a2a( core_attn_out: torch.Tensor, gate: torch.Tensor, thd_cp_a2a_inv: torch.Tensor | None, batch: int, seq_len: int, packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None = None, ) → torch.Tensor#

_apply_gated_norm(x, gate)#

_prepare_input_for_gated_delta_rule( qkv: torch.Tensor, gate: torch.Tensor, batch: int, seq_len: int, *gate_feats: tuple[torch.Tensor], ) → tuple[torch.Tensor, ...]#

Prepare the query, key, value, gate, and variant gate-feature tensors for the gated delta rule kernels.

Fuses split, reshape, L2 norm, repeat_interleave, and contiguous operations. gate_feats holds the variant-specific in_proj sections, which are returned contiguous for the decay/gating computation in forward.

_compute_g_and_beta( A_log_local_cp: torch.Tensor, dt_bias_local_cp: torch.Tensor, alpha: torch.Tensor, beta: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor]#: Compute g (decay) and beta (sigmoid) for gated delta rule. Fuses exp, softplus, mul, neg, and sigmoid operations.

_resolve_cu_seqlens( cu_seqlens_padded, cu_seqlens_actual, total_seq_len, name, cp_size: int = 1, ) → torch.Tensor#: Resolve cu_seqlens for packed sequence all-to-all, handling alignment padding.

sharded_state_dict( prefix='', sharded_offsets=(), metadata=None, tp_group=None, )#: Provide a sharded state dictionary for distributed checkpointing.

backward_dw()#: Execute weight gradient computation for all linear layers.

_backward_in_proj()#: Computes weight gradients of input projection layer.

_backward_out_proj()#: Computes weight gradients of output projection layer.

core.ssm.gated_delta_net.common._build_thd_cp_a2a_perm( cu_seqlens: torch.Tensor, cp_size: int, t_global: int, ) → tuple[torch.Tensor, torch.Tensor]#

core.ssm.gated_delta_net.common._build_head_perm_for_split_sections( split_sections: tuple[int, ...], cp_size: int, device: torch.device, ) → torch.Tensor#

core.ssm.gated_delta_net.common.get_parameter_local_cp( param: torch.Tensor, dim: int, cp_group: torch.distributed.ProcessGroup, split_sections: Optional[list[int]] = None, ) → torch.Tensor#

Get the local parameter for the current context parallel rank.

Parameters:

param (torch.Tensor) – The entire parameter to get the local parameter for.
dim (int) – The dimension to split the parameter along. Usually the dimension of head.
cp_group (torch.distributed.ProcessGroup) – The context parallel group.
split_sections (Optional[list[int]]) – If not None, first split the parameter along the dimension dim into sections, then get the local hidden parallel weights separately, finally concatenate the local hidden parallel weights along the dimension dim.

Returns:

The local parameter for the current context parallel rank.

Return type:

torch.Tensor

core.ssm.gated_delta_net.common.tensor_a2a_cp2hp( tensor: torch.Tensor, seq_dim: int, head_dim: int, cp_group: torch.distributed.ProcessGroup, split_sections: Optional[list[int]] = None, undo_attention_load_balancing: bool = True, )#

All-to-all context parallel to hidden parallel.

Parameters:

tensor (torch.Tensor) – The tensor to all-to-all. Currently only support (seq_len, batch, head_dim) shaped tensor.
seq_dim (int) – The dimension of sequence length. Currently only supports seq_dim == 0.
head_dim (int) – The dimension of head. Currently only supports head_dim == -1 or 2.
cp_group (torch.distributed.ProcessGroup) – The context parallel group.
split_sections (Optional[list[int]]) – If not None, split the tensor along the dimension head_dim into sections first, then do all-to-all for each section separately, finally concatenate the separated tensors along the dimension head_dim.
undo_attention_load_balancing (bool) – Whether to undo the attention load balancing of CP.

Returns:

The all-to-all tensor.

Return type:

torch.Tensor

core.ssm.gated_delta_net.common.tensor_a2a_hp2cp( tensor: torch.Tensor, seq_dim: int, head_dim: int, cp_group: torch.distributed.ProcessGroup, split_sections: Optional[list[int]] = None, redo_attention_load_balancing: bool = True, )#

All-to-all hidden parallel to context parallel.

Parameters:

tensor (torch.Tensor) – The tensor to all-to-all. Currently only support (seq_len, batch, head_dim) shaped tensor.
seq_dim (int) – The dimension of sequence length. Currently only supports seq_dim == 0.
head_dim (int) – The dimension of head. Currently only supports head_dim == -1 or 2.
cp_group (torch.distributed.ProcessGroup) – The context parallel group.
split_sections (Optional[list[int]]) – If not None, first split the tensor along the dimension head_dim into sections, then do all-to-all for each section separately, finally concatenate the separated tensors along the dimension head_dim.
redo_attention_load_balancing (bool) – Whether to redo the attention load balancing of HP.

Returns:

The all-to-all tensor.

Return type:

torch.Tensor

core.ssm.gated_delta_net.common.torch_chunk_gated_delta_rule( q, k, v, g, beta, chunk_size=64, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False, cu_seqlens=None, ) → tuple[torch.Tensor, torch.Tensor | None]#

Torch-native implementation of chunked gated delta rule for deterministic mode. Need this because FLA is not deterministic.

Reference: https://github.com/huggingface/transformers/blob/144c8ce2809a2e21914017652700e1ecb450501e/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L470-L547

core.ssm.gated_delta_net.common#

Module Contents#

Classes#

Functions#

Data#

API#

`core.ssm.gated_delta_net.common`#