`core.ssm.gated_delta_net.kda`#

Kimi Delta Attention, a channel-wise Gated DeltaNet variant.

Module Contents#

Classes#

`KimiDeltaAttentionSubmodules`	Submodules used by Kimi Delta Attention.
`KimiDeltaAttention`	Channel-wise Gated DeltaNet variant with direct Q/K/V/F/G projections.

API#

class core.ssm.gated_delta_net.kda.KimiDeltaAttentionSubmodules#

Bases: megatron.core.ssm.gated_delta_net.common.GatedDeltaNetSubmodules

Submodules used by Kimi Delta Attention.

beta_proj: megatron.core.transformer.spec_utils.ModuleSpec | type#: None

f_a_proj: megatron.core.transformer.spec_utils.ModuleSpec | type#: None

f_b_proj: megatron.core.transformer.spec_utils.ModuleSpec | type#: None

g_a_proj: megatron.core.transformer.spec_utils.ModuleSpec | type#: None

g_b_proj: megatron.core.transformer.spec_utils.ModuleSpec | type#: None

class core.ssm.gated_delta_net.kda.KimiDeltaAttention( config: megatron.core.ssm.kda_layer_config.KDALayerConfig, submodules: core.ssm.gated_delta_net.kda.KimiDeltaAttentionSubmodules, layer_number: int | None = 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 = None, *, name: str | None = None, cp_comm_type: str | None = None, pp_layer_offset: Optional[int] = None, is_mtp_layer: bool = False, )#

Bases: megatron.core.ssm.gated_delta_net.common._GDNBase

Channel-wise Gated DeltaNet variant with direct Q/K/V/F/G projections.

Supports direct or two-stage gates with equal query/key and value heads. Recurrent inference is not implemented.

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 (bool) – 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.
is_mtp_layer – Whether this module is inside an MTP prediction depth.

_setup_variant_attrs() → None#: Set KDA dimensions, projection checkpoint metadata, and kernel callable.

_get_feat_dim_split(cp_size_headwise: int) → tuple[int, int, int]#: Return KDA qkv/raw-g/output-gate split sizes for runtime headwise CP.

_reset_dt_bias() → None#: Initialize the KDA channel-wise step-size bias in inverse-softplus space.

_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]]#: Shape the raw channel-wise decay gate and activate the write strength.

_apply_gated_norm(x: torch.Tensor, gate: torch.Tensor) → torch.Tensor#: Apply per-head RMSNorm followed by KDA’s sigmoid output gate.

_proj(proj, x)#

Apply a projection GEMM, forcing BF16 when kda_disable_fp8 is set.

vLLM rollout keeps every KDA projection in BF16 (the checkpoint stores them BF16 with no FP8 scales), so the actor must match to keep actor<->rollout pearson sharp under FP8 hybrid training.

supports_two_stage_attention() → bool#: KDA currently implements only the atomic forward path.

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,

*,

pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None,

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

**kwargs,

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

Run KDA on inputs already arranged in config.linear_cp_layout.

HybridStack owns layout conversion through ContextParallelLayoutManager, including the packed metadata and output conversion at stack boundaries.

_forward_compute( hidden_states, batch, seq_len_post_headwise, cp_size_headwise, cp_group_headwise, cp_size_chunkwise, cp_group_chunkwise, cu_seqlens_q, packed_seq_params, chunkwise_cp_context, )#: Core KDA computation (in_proj -> conv1d -> gated_delta_rule -> norm -> out_proj).

_gated_norm_and_layout_restore( 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, cp_size_headwise: int, cp_group_headwise: torch.distributed.ProcessGroup | None, cp_size_chunkwise: int, cp_group_chunkwise: torch.distributed.ProcessGroup | None, cu_seqlens_q: torch.Tensor | None, ) → torch.Tensor#: Apply KDA’s gated output norm and restore its context-parallel layout.

pre_gated_delta_rule( qkvfg, beta, batch, seq_len, cp_size_headwise, cp_group_headwise, cu_seqlens_q=None, chunkwise_cp_context=None, packed_seq_params=None, *, raw_g=None, gate=None, )#: Prepare QKV, output gate, beta, and raw decay tensors before KDA.

static _validate_packed_cu_seqlens( cu_seqlens_q: torch.Tensor, cu_seqlens_kv: torch.Tensor, ) → None#: Validate the self-attention boundary contract for packed KDA.

backward_dw() → None#: Execute weight-gradient computation for all KDA projections.

core.ssm.gated_delta_net.kda#

Module Contents#

Classes#

API#

`core.ssm.gated_delta_net.kda`#