`core.ssm.kda_layer_config`#

Module Contents#

Configuration for a KDA (Kimi Delta Attention) layer in a hybrid stack.

class core.ssm.kda_layer_config.KDALayerConfig#

Bases: megatron.core.transformer.transformer_config.TransformerConfig

Configuration for a KDA (Kimi Delta Attention) layer in a hybrid stack.

linear_num_value_heads: Optional[int]#

KDA uses equal key and value head counts by default.

kda_disable_fp8: bool#

False

Force KDA projections to BF16 even under FP8 training, (KDA projections are BF16 in the checkpoint).

kda_safe_gate: bool#

False

Whether the KDA kernel should use bounded gate values.

kda_lower_bound: Optional[float]#

None

Optional lower bound for KDA’s bounded gate values.

kda_two_stage_gates: bool#

False

Use low-rank f_b(f_a(x)) and g_b(g_a(x)) gates with a QKV-only input projection.

validate_kda() → None#: Validate KDA-specific dimensions, including configs copied with from_config.