[HollowMan6]

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

Fixes #7983

What does this PR do?

Add KDA (Kimi Delta Attention) hybrid layers to HybridModel, including the two-stage low-rank gates used by GLM-5.3-Flash, and an option to keep the KDA projections in BF16 under FP8 training.

This PR was split into four, as discussed with @Phlip79 in the comments. It now carries only the KDA layers; the other parts live in their own PRs, each independent of the others and based on main:

• mHC precision controls, mean output contract, full activation recompute with mHC: Add mHC precision controls, mean output contract, and full activation recompute with mHC #7521
• KPool DSA indexer, NoPE for the DSA indexer / absorbed MLA, unquantized indexer and MLA attention GEMMs under FP8: Add KPool DSA indexer, NoPE MLA support, and FP8 precision controls for the indexer and MLA attention #7522
• Blockwise FP8 backward (E5M2) wgrad guard in fp8_utils.py: Guard blockwise FP8 backward (E5M2) wgrad against all-zero gradient blocks #7523

GLM-5.3-Flash at a glance (from its HF config.json): 45 layers = 34 × KDA (linear_attention) + 11 × DSA (deepseek_sparse_attention) over NoPE MLA; mHC with 4 residual streams; KDA gate_lower_bound = -5.0; 1 MTP depth; 288 routed experts. KDA layers use the new K symbol in the hybrid layer pattern.

Changes

• Ported from dev: KimiDeltaAttention in ssm/gated_delta_net/kda.py (channel-wise gated DeltaNet on FLA chunk_kda, separate beta_proj, headwise / chunkwise CP, packed THD), resolve_cp_group and cp_partition_mode in packed_seq_params.py, nvtx_range in utils.py. CP layout conversion goes through HybridStack's ContextParallelLayoutManager (megatron/core/context_parallel/layout.py); the earlier context_parallel_layout/ port is gone.
• Added on top of dev's KDA:
  • kda_two_stage_gates: GLM-style low-rank gates f_b(f_a(x)) (decay) and g_b(g_a(x)) (output) with a QKV-only in_proj. The decay gate is precomputed in FP32 with FLA fused_kda_gate (honouring kda_safe_gate / kda_lower_bound), and the output RMSNorm + sigmoid gate run in FLA rms_norm_gated, so there is a single BF16 rounding after norm and gating. f_a / g_a are replicated (TELinear, parallel_mode="duplicated"); f_b / g_b are column-parallel.
  • A_log / dt_bias are FP32 and marked keep-in-FP32, so they survive Float16Module, optimizer construction and checkpointing.
  • kda_disable_fp8: run every KDA projection GEMM with FP8 disabled (see Motivation).
  • Selective recompute of the whole KDA block with recompute_modules=["gdn"] (gdn_norm_out keeps working).
• Hybrid wiring: KDALayerConfig, layer symbol K, kda_layer slot in HybridStackSubmodules, layer specs in hybrid_stack_spec and hybrid_inference_stack_spec.
• _GDNBase (common.py): in_proj sizing for two-stage gates, per-variant gate-parameter dtype, headwise CP size override in _prepare_input_for_gated_delta_rule, chunkwise CP-context cache, is_mtp_layer.
• HybridModel.forward: the same output_processor / output_processor_context hook as GPTModel (caller-owned output projection and loss).
• New TransformerConfig fields: kda_two_stage_gates, kda_safe_gate, kda_lower_bound, kda_disable_fp8, gdn_conv_pad_alignment, gdn_pre_gated_delta_rule_fusion (reserved; raises NotImplementedError for KDA).

Motivation

With an FP8 actor and a vLLM rollout, every module the rollout computes in BF16 has to be computed in BF16 on the actor as well, otherwise the two diverge. The released FP8 checkpoint keeps every KDA projection (q/k/v_proj, b_proj, f_* / g_*_proj, o_proj) in BF16 (quantization_config.modules_to_not_convert); kda_disable_fp8 lets the actor reproduce that split while the MoE experts (the bulk of the FLOPs) stay in FP8.

Relationship to dev and other PRs

• From dev with additions: kda.py, the kda_layer slot in hybrid_block.py, resolve_cp_group, nvtx_range.
• New here (not on dev): two-stage KDA gates, kda_disable_fp8, KDALayerConfig / K symbol, HybridModel.output_processor.
• Related: [dev] feat(moe): Add KDA Low-Rank Projections and Extend MLA Support #6877 (dev: KDA low-rank projections and MLA NoPE), Add Kimi Delta Attention (KDA) linear attention variant #5769 (community KDA as an attention variant), Feature Request: Kimi Delta Attention (KDA) #2446 (KDA feature request), GLM-5.2 training support tracking #6392 (GLM-5.2 training support tracking).

Known gaps

• KDA context parallelism: the headwise / chunkwise CP paths and cp_partition_mode are included, but only CP = 1 has been validated in this PR.
• KDA has no recurrent inference path (NotImplementedError), and gdn_pre_gated_delta_rule_fusion is not implemented for KDA.
• KDA dispatches to FLA kernels (chunk_kda, fused_kda_gate, rms_norm_gated) and adds jit_fuser helpers; there is no entry in the kernel determinism manifest yet.
• No functional test / golden values yet. The change touches numerics, so Run functional tests is the right CI label.

Testing

• tests/unit_tests/ssm/test_kda_gate_precision.py (new): FP32 gate parameters survive BF16 wrapping, gated-norm numerics, packed-THD output equals per-sequence output, for kda, kda_direct and gdn.
• tests/unit_tests/ssm/test_hybrid_block.py: the K symbol builds the kda_layer slot and selects the linear-attention CP layout.

pytest tests/unit_tests/ssm/test_kda_gate_precision.py tests/unit_tests/ssm/test_hybrid_block.py

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue: Related to #2446, #6392

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 (guide)
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

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

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[Copilot]

🟢 Approval recommended

The logic change is small, follows existing get_fp8_disabled_context usage patterns in related indexer code, and the only feedback is minor comment wording accuracy.

Pull request overview

Adjusts DSA indexer precision behavior under hybrid FP8 training to reduce quantization noise in the top‑k index score path, improving train/inference alignment for sparse-attention selection.

Changes:

• Disable FP8 quantization for linear_wk and linear_weights_proj during module construction so their parameters stay in higher precision when FP8 params are enabled.
• Disable FP8 quantization for the K-path (linear_wk, k_norm, RoPE, optional activation rotation) and linear_weights_proj during forward, while keeping the Q-path (linear_wq_b) quantized.

File summaries

File	Description
megatron/core/transformer/experimental_attention_variant/dsa.py	Split DSA indexer Q vs K/weights-projection execution so index-score-critical ops run with FP8 disabled.

Review details

Suppressed comments (1)

megatron/core/transformer/experimental_attention_variant/dsa.py:1401

• The updated comments call out “(FP8)” and “run in BF16”, but get_fp8_disabled_context only disables TE quantization and doesn’t force a specific dtype. Rewording to “quantized” vs “FP8/FP4 disabled” avoids implying these paths are always FP8/BF16 regardless of the model’s base precision configuration.

        # q linear and apply rope to q  (FP8)
        # =========================================
        # [seqlen, batch, q_lora_rank] -> [seqlen, batch, index_n_heads * index_head_dim]
        q, _ = self.linear_wq_b(qr)
        # [seqlen, batch, index_n_heads * index_head_dim]
        #   -> [seqlen, batch, index_n_heads, index_head_dim]
        q = q.reshape(seqlen, bsz, self.index_n_heads, self.index_head_dim)
        q = self._apply_rope(q, rotary_pos_emb, mscale, cu_seqlens=cu_seqlens_q)
        if self.config.dsa_indexer_rotate_activation:
            q = rotate_activation(q)

        # =========================================
        # k linear, k_norm, rotate, and weights_proj run in BF16 (FP8 disabled).
        # =========================================

• Files reviewed: 1/1 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[erictang000]

thanks for this PR @HollowMan6, also testing this branch out now for RL support on SkyRL

[HollowMan6]

@Phlip79 I validated GLM 5.3 flash RL with LoRA together with vLLM/Megatron-Bridge/verl using geo3k, you can find the convergence results at vllm-project/vllm#56327 (comment)

[cuichenx]

/ok to test 22c4f37

[Phlip79]

/ok to test f8248c0

[kevalmorabia97]

LGTM from the Model Optimizer side. Tested 92bd36ec (merged into #7521) and NVIDIA-NeMo/Megatron-Bridge#6044 in a NeMo 26.10 container on GB300: ModelOpt GPTQ PTQ, distillation with save + resume, and quantized HF export on a 4-layer GLM-5.3-Flash, plus full-size GLM-5.3-Flash QAD on 32 nodes (EP32 x CP4, packed SFT) with distributed-optimizer save + resume, matching the loss and grad norms of our earlier full PTQ + QAD + vLLM eval study on older heads of this stack.

NVIDIA-NeMo/Megatron-Bridge#6044 and our GLM-5.3-Flash support (NVIDIA/Model-Optimizer#2539) are waiting on this PR, so we'd appreciate a merge when possible. Thanks @HollowMan6!

[kevalmorabia97]

/ok to test 177c4f1

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40762445544