core.ssm.gated_delta_net.gdn#
Module Contents#
Classes#
Functions#
Get the local parameter slice for headwise context parallelism. |
|
Torch-native implementation of chunked gated delta rule for deterministic mode. Need this because FLA is not deterministic. |
API#
- core.ssm.gated_delta_net.gdn.get_parameter_local_cp_headwise(
- param: torch.Tensor,
- dim: int,
- cp_group: Optional[torch.distributed.ProcessGroup],
- split_sections: Optional[list[int]] = None,
Get the local parameter slice for headwise context parallelism.
- 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.
- _get_feat_dim_split(cp_size: int) tuple[int, int, int, int]#
Return projection split sizes for the runtime head-parallel CP size.
- _compute_gates(
- A_log_local_cp: torch.Tensor,
- dt_bias_local_cp: torch.Tensor,
- batch: int,
- seq_len: int,
- *gate_feats: tuple[torch.Tensor],
Compute the per-head log-decay g and the write strength beta.
- _apply_gated_norm(x: torch.Tensor, gate: torch.Tensor) torch.Tensor#
Normalize and gate the recurrence output, validating every fused call.
- _apply_gated_norm_unfused(x, gate)#
Preserve the projection-backed gate view in the unfused output path.
- 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,
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,
Dispatch inference, then fall through to the training two-stage path.
ssm_dynamic_inferenceapplies the output projection itself, so its result already satisfies this method’s(output, bias)contract and is returned directly. It cannot live inforward_pre_attn_and_core_attn, whose contract is the tensor before that projection, because the baseforwardwould 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,
- cp_size: int,
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,
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,
Run one CUDA-graph-compatible GDN decode token per request.
Both kernels read and write the state caches in place at
batch_indices; a-1(padding) row gets a zero output without reading inputs or state.
- ssm_prefill(
- projected: torch.Tensor,
- conv_state: torch.Tensor,
- ssm_state: torch.Tensor,
- context: megatron.core.inference.contexts.DynamicInferenceContext,
Run graph-safe GDN prefill using the single-update GDP recurrence.
Chunk descriptors belong to the inference context and are refreshed in place before replay. Their padded capacity fixes every workspace and launch grid independently of the lengths used during graph capture.
- pre_gated_delta_rule(
- qkvzba,
- batch,
- seq_len,
- cp_size,
- cp_group,
- cu_seqlens_q=None,
Prepare QKV, gate, beta, and decay tensors before the gated delta rule.
- _fused_streamed_pre_gated_delta_rule(
- qkvzba,
- cu_seqlens_q=None,
- seq_idx=None,
- cp_size_headwise=1,
- cp_group_headwise=None,
Call the streamed fused pre-GDR wrapper.
- 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,
Torch-native implementation of chunked gated delta rule for deterministic mode. Need this because FLA is not deterministic.
scaledefaults to1 / 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