`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 linear attention interfaces, compliant to upstream FLA 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.
`a2a_cp_to_hp`	Run GDN context-parallel to hidden-parallel A2A and return its inverse context.
`a2a_hp_to_cp`	Run GDN hidden-parallel to context-parallel A2A using CP-to-HP context.

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 linear attention interfaces, compliant to upstream FLA interfaces.

Only q/k/v/g are common to every kernel, and only as keywords: each variant inserts its own gates after g (e.g., beta for GDN, b/w for GDN2).

__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, *, is_mtp_layer: bool = False, name: str | None = None, cp_comm_type: str | None = None, pp_layer_offset: int = 0, )#

Bases: megatron.core.transformer.module.MegatronModule, megatron.core.transformer.module.TwoStageAttentionLayer

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.
is_mtp_layer – Accepted for TransformerLayer compatibility and ignored because GDN does not have MTP-specific construction behavior.
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.
pp_layer_offset – Offset of this pipeline stage’s first global layer.

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

supports_two_stage_attention() → bool#: Output-norm recomputation requires the original atomic forward path.

forward_post_core_attn( norm_out: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor | None]#: Apply a GDN variant’s output projection to its normalized recurrence output.

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_dt_bias()#

Initialize dt_bias. Called from reset_parameters under the RNG tracker.

Defaults to ones; variants whose kernel expects a different step-size parametrization override this.

reset_parameters()#: Reset the parameters.

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 | None]#: Run a GDN variant’s recurrence followed by its output projection.

_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, A_log_local_cp: torch.Tensor, dt_bias_local_cp: torch.Tensor, batch: int, seq_len: int, *gate_feats: tuple[torch.Tensor], ) → dict[str, torch.Tensor]#

Prepare all gated delta rule kernel inputs.

Fuses split, reshape, L2 norm, decay/gate activations, repeat_interleave, and contiguous operations. gate_feats holds the variant-specific in_proj sections, which _compute_gates turns into the decay and gating tensors.

Returns:: Kernel inputs keyed by kernel argument name (q, k, v, g, plus the variant-specific gates), and the output gate (z) tensor under the gate key, which is not a kernel input.
Return type:: (dict[str, Tensor])

abstractmethod _compute_gates( A_log_local_cp: torch.Tensor, dt_bias_local_cp: torch.Tensor, batch: int, seq_len: int, *gate_feats: tuple[torch.Tensor], ) → tuple[torch.Tensor, dict[str, torch.Tensor]]#

Compute the log-decay g and the variant-specific kernel inputs.

Parameters:

A_log_local_cp – CP-local slice of A_log.
dt_bias_local_cp – CP-local slice of dt_bias.
batch – Batch size.
seq_len – Sequence length.
gate_feats – The variant-specific in_proj output sections (everything after the qkv and output-gate sections, in feat_dim_split order).

Returns:

The log-decay g and a dict of the remaining variant-specific kernel inputs keyed by kernel argument name.

Return type:

(tuple[Tensor, dict[str, Tensor]])

_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.a2a_cp_to_hp( qkvzba: torch.Tensor, in_proj_split_sections: tuple[int, ...], cp_size: int, cp_group: torch.distributed.ProcessGroup, cu_seqlens_q: torch.Tensor | None, seq_len: int, packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None, ) → tuple[torch.Tensor, torch.Tensor | None]#

Run GDN context-parallel to hidden-parallel A2A and return its inverse context.

Parameters:

qkvzba – in_proj output in sbhd format, sharded along the sequence dim over CP.
in_proj_split_sections – per-section sizes of the in_proj output, local to this TP rank, used to build the pre-a2a head permutation.
cp_size – context-parallel world size.
cp_group – context-parallel process group.
cu_seqlens_q – cumulative sequence lengths, required for the thd path.
seq_len – global (unsharded) sequence length.
packed_seq_params – packed-sequence params; the thd path is taken when its qkv_format is 'thd'.

Returns:

The hidden-parallel tensor and the sequence-dim inverse permutation to hand to

func:: a2a_hp_to_cp (None outside the thd + CP>1 case).

core.ssm.gated_delta_net.common.a2a_hp_to_cp( norm_out: torch.Tensor, cp_size: int, cp_group: torch.distributed.ProcessGroup, packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None, thd_cp_a2a_inv: torch.Tensor | None, ) → torch.Tensor#

Run GDN hidden-parallel to context-parallel A2A using CP-to-HP context.

Parameters:

norm_out – gated-norm output in sbhd format, sharded along the head dim over CP.
cp_size – context-parallel world size.
cp_group – context-parallel process group.
packed_seq_params – packed-sequence params; the thd path is taken when its qkv_format is 'thd'.
thd_cp_a2a_inv –
sequence-dim inverse permutation returned by

func:

a2a_cp_to_hp, required on the thd path when cp_size > 1.

Returns:

The context-parallel tensor, matching the layout of the GDN module input.

core.ssm.gated_delta_net.common#

Module Contents#

Classes#

Functions#

Data#

API#

`core.ssm.gated_delta_net.common`#