feat: Kimi Delta Attention (KDA) attention variant, and configurable mHC input-norm epsilon

[erictang000]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds Kimi Delta Attention (KDA) as experimental_attention_variant="kda", and makes the mHC input-normalization epsilon placement configurable. Together these are the two megatron-core pieces needed to build GLM-5.3-Flash (HF glm5_next) out of mcore components; the companion NVIDIA-NeMo/Megatron-Bridge#6137 adds the provider/spec/bridge on top and is blocked on this PR.

Issue tracking

Linked issue: not yet filed — happy to open a feature request if the team wants one before review.

The two changes

1. feat(mhc): configurable input-normalization epsilon placement

HyperConnectionModule normalizes the flattened residual streams as x / (rms(x) + eps) with eps hard-coded to 1e-6. Models trained with a standard RMSNorm in that position — x * rsqrt(mean(x^2) + eps) — cannot be loaded into mHC as it stands.

The two forms agree for O(1) activations but not for small residual streams. GLM-5.3-Flash uses rms_norm_eps (1e-5) inside the square root and its embeddings have a per-token rms below sqrt(1e-5), where the placement of the epsilon changes the mixing weights materially — not a tolerance that can be absorbed.

Two TransformerConfig fields:

field	default	meaning
mhc_norm_eps	1e-6	the epsilon itself, previously a literal in __init__
mhc_norm_eps_inside_sqrt	False	selects the standard-RMSNorm form, via a new native_proj_rms_eps_inside_sqrt op

Both defaults reproduce today's behavior exactly. The fused mHC kernels implement the paper form only, so use_fused_mhc together with mhc_norm_eps_inside_sqrt=True is rejected in __post_init__ rather than silently ignored.

2. feat(ssm): KimiDeltaAttention

KDA (Kimi Linear, arXiv:2510.26692) is a gated delta rule whose forget gate is per channel rather than per head:

q, k, v = SiLU(conv1d(W_q x)), SiLU(conv1d(W_k x)), SiLU(conv1d(W_v x))   # depthwise, causal
f       = W_fb (W_fa x)                                                   # low-rank forget gate
g       = lower_bound * sigmoid(exp(A_log) * (f + dt_bias))               # log-decay in [lb, 0)
beta    = sigmoid(W_b x)
o       = chunk_kda(l2norm(q), l2norm(k), v, g, beta)                     # fla kernel
y       = W_o (RMSNorm(o) * sigmoid(W_gb (W_ga x)))                       # gated output norm

megatron/core/ssm/kda.py adds the module next to the gated delta net family; is_linear_attention_variant now covers "kda", so the existing linear_attention_freq pattern and the hybrid block builder drive it unchanged.

• Unfused projections. Unlike GDN, KDA keeps its projections separate rather than fusing them into one in_proj. That is the layout its reference checkpoints ship in, and keeping it means an HF bridge maps 1:1. Consequently fuse_input_layernorm=False. Open to folding this onto _GDNBase instead if you'd prefer one code path — it would cost the checkpoint-layout correspondence.
• TP shards heads. q/k/v/f_b/g_b/b column-parallel; the three depthwise causal convolutions and the fp32 A_log/dt_bias split along the head dimension; the low-rank f_a/g_a down-projections duplicated; o_proj row-parallel. The replicated o_norm weight is marked sequence_parallel so finalize_model_grads sums its gradient across the TP group.
• Packed sequences (qkv_format="thd") go through cu_seqlens. sharded_state_dict follows GDN.
• Config. Reuses the shared linear_* fields — KDA requires the q/k and v head counts and head dimensions to be equal, enforced in __post_init__ — plus a new kda_gate_lower_bound. Setting it bounds the log-decay to [lower_bound, 0) (GLM-5.3-Flash uses -5.0); the default None keeps the unbounded -exp(A_log) * softplus(f + dt_bias) gate of the original Kimi Linear.
• Not supported: context parallelism and inference caches. Both raise rather than degrade.

Validation

Run against this branch on a single B300, torch 2.13, TE from the same toolchain:

suite	result
tests/unit_tests/ssm/test_kda.py (new)	1 passed — packed two-sequence forward vs. the HF glm5_next reference layer in bf16
tests/unit_tests/transformer/test_hyper_connection_norm_eps.py (new)	4 passed — both norm forms vs. their closed-form reference, the config guard, and the small-scale regime that motivates the knob
tests/unit_tests/transformer/test_hyper_connection_recompute.py + test_mhc_block_manager.py (existing)	30 passed — no regression from the hyper_connection.py change

Both modules were also validated end-to-end in their out-of-tree form (same code, carried against b3393bbb) as part of a GLM-5.3-Flash bring-up: Megatron vs. HF transformers logits on the 4-layer GLM-5.3-Flash slice came out at the bf16 noise floor (KL 0.0082, 91.5% argmax at TP1, against a measured HF-bf16-vs-fp32 floor of 0.0071 / 95.7%), and a vLLM↔Megatron logprob round-trip on 4×H100 at TP2/EP4 matched to 0.064 mean absolute logprob. Details: NovaSky-AI/SkyRL#2179.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py — please advise: native_proj_rms_eps_inside_sqrt is a new torch.compile op and KDA dispatches to fla's chunk_kda, so this likely applies; I did not want to guess at the manifest entries.
• I have added proper typing to my code
• I have added relevant documentation (config docstrings; module docstring in kda.py)
• I have run the autoformatter on my PR (black 26.3.0 / isort 5.13.2; pylint 10.00/10 on the changed megatron/core files)

Known gaps

• KDA has no context-parallel path and no inference cache — both raise.
• k-pool DSA indexer support and hybrid-aware DSA index sharing are the remaining megatron-core gaps for GLM-5.3-Flash; they are not in this PR.

🤖 Generated with Claude Code

[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.

[erictang000]

close in favor of #7054