Add Kimi Delta Attention (KDA) linear attention variant

[CodersAcademy006]

What

Adds Kimi Delta Attention (KDA) as a new experimental_attention_variant, implementing the algorithm from Kimi Linear: An Expressive, Efficient Attention Architecture (https://arxiv.org/abs/2510.26692). This is the clean, single-purpose replacement requested in #2446, superseding the earlier #2573 and #2674 (which did not implement the paper's KDA algorithm).

The algorithm

KDA extends the gated delta rule with a fine-grained, per-key-channel forget gate Diag(alpha_t) (a channel-wise vector in [0,1]^{d_k}) in place of GatedDeltaNet's scalar per-head decay (Eq. 1 of the paper):

S_t = (I - beta_t k_t k_t^T) Diag(alpha_t) S_{t-1} + beta_t k_t v_t^T
o_t = S_t^T q_t

Following the maintainer guidance in #2446, the official Moonshot design and the flash-linear-attention chunk_kda operator are used as the numerical reference. The recurrence, chunk-wise parallel form, q/k L2-normalization, beta sigmoid and the -exp(A_log) * softplus(f_proj(x) + dt_bias) decay are all evaluated inside chunk_kda.

What this PR adds

• megatron/core/ssm/kimi_delta_attention.py — KimiDeltaAttention + KimiDeltaAttentionSubmodules, mirroring the structure of gated_delta_net.py. The module owns the q/k/v/beta projections, the two low-rank bottlenecks (f_proj for the fine-grained gate, g_proj for the output gate), the depthwise short convolution, the decay parameters (A_log, dt_bias) and the gated output RMSNorm.
• TransformerConfig — new kimi_delta_attention value for experimental_attention_variant, sharing the linear-attention validation with gated_delta_net.
• experimental_attention_variant_module_specs.py — get_kimi_delta_attention_module_spec, dispatch, and registration as a linear-attention variant (so it follows the linear_attention_freq LA/SDPA layer pattern, matching the paper's 3:1 KDA-to-global ratio).
• Unit test tests/unit_tests/ssm/test_kimi_delta_attention.py (forward + sharded_state_dict, TP/SP parametrized, FLA-gated).

Design notes

• Separate q/k/v/beta projections (rather than GDN's fused input projection), matching the FLA reference and keeping tensor-parallel head sharding straightforward.
• Low-rank f_proj/g_proj built with the established MLA/DSA pattern in this file: replicated (duplicated) down-projection, column-parallel up-projection.
• fuse_input_layernorm=False so every projection consumes the same normalized hidden states (KDA has multiple projection groups, unlike GDN's single fused input projection).
• GVA: q/k stay at num_key_heads and v/gate/beta at num_value_heads; grouping is handled inside chunk_kda (no repeat_interleave).
• Tensor-parallel sharded over heads, with a sharded_state_dict for distributed checkpointing.

Scope (this draft)

Focused on a correct, reviewable core. Context parallelism, packed (THD) sequences and inference/KV-cache are explicitly guarded with NotImplementedError for now and can be layered on following the GDN CP/THD machinery in a follow-up.

Testing

• Tensor plumbing validated on CPU across TP=1/2/4: q/k/v/g/beta/A_log/dt_bias shapes all satisfy the chunk_kda contract and the output is [s, b, h].
• Unit test added (forward + sharded checkpoint). black/isort clean.
• The chunk_kda kernel requires CUDA + flash-linear-attention, so numerical correctness and TP/SP execution are left to CI / GPU (I don't have the hardware to run the FLA kernel or the TE-backed unit test locally). Happy to iterate on anything CI surfaces.

Requires a flash-linear-attention build that exposes fla.ops.kda.chunk_kda.

Refs #2446. cc @sbhavani @BoxiangW — this is the clean replacement you asked for; happy to iterate on scope (CP/THD/inference) and on any fusion you'd prefer for the projections.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[sbhavani]

@CodersAcademy006 thanks for the efforts! we're adding KDA here moving forward: #6877