Is this page helpful?

core.ssm.gated_delta_net.gdn#

Module Contents#

Classes#

GatedDeltaNet

Functions#

torch_chunk_gated_delta_rule

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

API#

class core.ssm.gated_delta_net.gdn.GatedDeltaNet(
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.ssm_inference.SSMDynamicInferenceMixin, megatron.core.ssm.gated_delta_net.common._GDNBase

_setup_variant_attrs()#

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

_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-head log-decay g and the write strength beta.

forward_pre_attn_and_core_attn(
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,
packed_sequence_cp_metadata=None,
**kwargs,
) → torch.Tensor#

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

Returns:

Normalized recurrence output.

Return type:

torch.Tensor

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]#

Dispatch inference, then fall through to the training two-stage path.

ssm_dynamic_inference applies the output projection itself, so its result already satisfies this method’s (output, bias) contract and is returned directly. It cannot live in forward_pre_attn_and_core_attn, whose contract is the tensor before that projection, because the base forward would then project it a second time.

Returns:

GDN output and bias.

Return type:

tuple[torch.Tensor, torch.Tensor | None]

_split_projection(
projected: torch.Tensor,
batch: int,
seq_len: int,
) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]#

Split the fused projection into qkv, output gate, beta, and alpha.

_prepare_inference_inputs(
qkv: torch.Tensor,
beta: torch.Tensor,
alpha: torch.Tensor,
batch: int,
seq_len: int,
) → dict[str, torch.Tensor]#

Prepare raw FLA inputs while leaving normalization and gates fused in-kernel.

mamba_state_shapes_per_request() → tuple[tuple[int, ...], tuple[int, ...]]#

Return the TP-local convolution and delta-rule cache shapes.

ssm_decode(
projected: torch.Tensor,
conv_state: torch.Tensor,
ssm_state: torch.Tensor,
batch_indices: torch.Tensor,
intermediate_conv_state: torch.Tensor | None = None,
intermediate_ssm_state: torch.Tensor | None = None,
) → torch.Tensor#

Run one CUDA-graph-compatible GDN decode token per request.

ssm_prefill(
projected: torch.Tensor,
conv_state: torch.Tensor,
ssm_state: torch.Tensor,
context: megatron.core.inference.contexts.DynamicInferenceContext,
) → torch.Tensor#

Run packed variable-length GDN prefill and populate request states.

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

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

scale defaults to 1 / sqrt(K), matching the FLA kernel. Extra keyword arguments are accepted and ignored so this stays interchangeable with the FLA kernel, which takes several options this implementation does not model.

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