Is this page helpful?

core.ssm.gated_delta_net.gdn2#

Module Contents#

Classes#

GatedDeltaNet2

GDN2 (Gated DeltaNet-2) layer class.

Functions#

torch_chunk_gdn2

Torch-native chunkwise Gated Delta Rule-2, for deterministic mode.

Data#

logger

API#

core.ssm.gated_delta_net.gdn2.logger#

‘getLogger(…)’

class core.ssm.gated_delta_net.gdn2.GatedDeltaNet2(
config: megatron.core.transformer.TransformerConfig,
submodules: megatron.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.ssm.gated_delta_net.common._GDNBase

GDN2 (Gated DeltaNet-2) layer class.

GDN2 replaces GDN’s per-head scalar decay and write strength with channel-wise gates, decoupling erase and write:

S_t = (I - k_t (b_t * k_t)^T) Diag(exp(g_t)) S_{t-1} + k_t (w_t * v_t)^T

where g_t is a per-key-channel log-decay, b_t (in R^{d_k}) is the channel-wise erase gate, and w_t (in R^{d_v}) is the channel-wise write gate. Reference: “Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention” (https://github.com/NVlabs/GatedDeltaNet-2).

Note: unlike the GDN2 reference implementation, which uses low-rank decay and output-gate projections, all GDN2 projections are fused full-rank into the single column-parallel in_proj for TP/CP/SP simplicity.

The layer takes input with size [s, b, h] and returns output of the same size.

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.

_setup_variant_attrs()#

Set the GDN2 in_proj sizing, split tables, gate parameter dims, and kernel.

_reset_dt_bias()#

Softplus-inverse init of dt_bias.

Initializes so the initial per-channel step size lands in [1e-3, 0.1], following the GDN2 reference implementation.

_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 per-channel log-decay g and the erase/write gates b/w.

forward_pre_attn_and_core_attn(
hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
inference_context: megatron.core.inference.contexts.BaseInferenceContext | None = None,
packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None = None,
sequence_len_offset: int | None = None,
*,
inference_params: megatron.core.inference.contexts.BaseInferenceContext | None = None,
packed_sequence_cp_metadata=None,
**kwargs,
) → torch.Tensor#

Run GDN2 through its normalized recurrence output, before output projection.

Returns:

Normalized recurrence output.

Return type:

torch.Tensor

core.ssm.gated_delta_net.gdn2.torch_chunk_gdn2(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
b: torch.Tensor,
w: torch.Tensor,
scale: float | None = None,
chunk_size: int = 64,
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]#

Torch-native chunkwise Gated Delta Rule-2, for deterministic mode.

Parameters:

  • q – queries of shape [B, T, H, K].

  • k – keys of shape [B, T, H, K].

  • v – values of shape [B, T, H, V].

  • g – channel-wise log-decay of shape [B, T, H, K].

  • b – channel-wise erase gate of shape [B, T, H, K].

  • w – channel-wise write gate of shape [B, T, H, V].

  • scale – attention scale. Defaults to 1 / sqrt(K).

  • chunk_size – chunk length of the WY schedule.

  • initial_state – optional [B, H, K, V] initial state.

  • output_final_state – whether to also return the final recurrent state.

  • use_qk_l2norm_in_kernel – L2-normalize q and k here rather than in the caller.

  • cu_seqlens – packed-sequence offsets; unsupported, must be None.

  • kwargs – accepted and ignored, so this stays interchangeable with the FLA kernel, which takes several options this implementation does not model.

Returns:

output of shape [B, T, H, V] and the final state, or None when output_final_state is False.

Return type:

(tuple[Tensor, Tensor | None])