GLM5.3 Flash (KDA + mHC + KPool DSA) support + FP8

[HollowMan6]

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

What does this PR do?

Add the model-side pieces needed to train GLM-5.3-Flash on Megatron Core's HybridModel, plus the precision controls needed to keep an FP8-trained actor numerically aligned with a vLLM rollout of the released checkpoint:

1. KDA (Kimi Delta Attention) hybrid layers, including GLM's two-stage low-rank gates.
2. mHC (Manifold-Constrained Hyper-Connections) numerics used by GLM, and full activation recompute with mHC.
3. KPool DSA indexer (pooled sparse-attention index selection) and NoPE for the DSA indexer / absorbed MLA.
4. FP8 precision splits for the DSA indexer, the MLA attention GEMMs and the KDA projections.

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 (qk_rope_head_dim = 0); mHC with 4 residual streams and 20 Sinkhorn iterations; indexer with index_kpool = 4, index_kpool_always_select_tail = true, index_topk = 2048; 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

1. KDA in HybridModel (commit KDA)

• Ported from dev: megatron/core/context_parallel_layout/ (zigzag ↔ contiguous CP layout conversion for SBHD and THD; byte-identical to dev), resolve_cp_group in packed_seq_params.py, nvtx_range in utils.py, and KimiDeltaAttention in ssm/gated_delta_net/kda.py (channel-wise gated DeltaNet on FLA chunk_kda, separate beta_proj, headwise/chunkwise CP, packed THD).
• 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 §4).
  • 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).

2. mHC (commit mHC)

• mhc_norm_eps_inside_sqrt: compute the mapping norm as rsqrt(mean(x²) + layernorm_epsilon) instead of ‖x‖ / √K + 1e-6.
• mhc_keep_mappings_in_fp32: keep h_pre / h_post / h_res and the residual-stream mixing in FP32 until the final cast to the activation dtype.
• mhc_learned_output_contract (default True, the current behaviour): with False, the residual streams are contracted by their mean instead of the learned hc_head_* weights, in both the HybridStack post-process and the MTP layer (GLM-5.3-Flash has no learned head).
• Full activation recompute with mHC-wrapped hybrid layers: HyperConnectionHybridLayer sets supports_hybrid_recompute_kwargs, recompute.py forwards the matching kwargs (including packed_sequence_cp_metadata), and the NotImplementedError in TransformerConfig.__post_init__ is removed.
• use_fused_mhc raises ValueError together with either new precision flag (the fused kernels do not implement them).

3. KPool DSA indexer + NoPE (commit KPool DSA)

• dsa_indexer_kpool > 1: keys are compressed kpool tokens at a time with a per-token compression gate (index_kpool_compress_gate, BF16, same parameter name as the checkpoint) plus a per-slot additive position bias (index_kpool_compress_ape, FP32): per-dimension softmax over the slots, FP32 accumulation, BF16 pooled key. Top-k runs over pools (index_topk // kpool) and is expanded back to token indices (fused_qk_topk_kpool); dsa_indexer_kpool_always_select_tail appends each query's incomplete causal pool so short prefixes stay fully visible. Packed THD restarts pools at document boundaries. The fused cuDNN DSA path is bypassed for KPool, and the head-weight projection stays in FP32.
• dsa_indexer_kpool_fp8: emulate the serving indexer's FP8 input path (FP32 Hadamard → BF16 → E4M3 with a power-of-two row scale) on q and on the pooled keys, so training top-k matches inference.
• NoPE: qk_pos_emb_head_dim == 0 is supported by DSAIndexer and AbsorbedMLA (no rotary embedding is built; the absorbed-q path has no positional slice).
• mla_disable_attention_fp8: run linear_q_down_proj, linear_q_up_proj, linear_kv_down_proj and linear_proj of AbsorbedMLA with FP8 disabled.

4. DSA indexer FP8 precision split (commits fix: DSA indexer FP8 precision split … and linear_wq_b also in bf16)

linear_wk and linear_weights_proj are built under get_fp8_disabled_context(config, is_init=True) (BF16 parameters, no FP8 scales), and in forward_before_topk the k path (linear_wk, k_norm, RoPE, activation rotation), linear_weights_proj and the linear_wq_b GEMM all run under get_fp8_disabled_context. Net effect: the whole indexer runs unquantized while the rest of the model keeps FP8.

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 the whole DSA indexer (wq_b, wk, weights_proj, k_norm, index_kpool_compress_*) and every KDA projection (q/k/v_proj, b_proj, f_* / g_*_proj, o_proj) in BF16 (quantization_config.modules_to_not_convert), and vLLM runs the main MLA attention path in BF16. The indexer matters most: it feeds a discrete top-k, so FP8 noise there changes which tokens are attended, and actor/rollout agreement (Pearson) degrades. These options let the actor reproduce the serving precision split while the MoE experts (the bulk of the FLOPs) stay in FP8.

Relationship to dev and other PRs

• Byte-identical to dev: megatron/core/context_parallel_layout/*.
• From dev with additions: kda.py, the kda_layer slot in hybrid_block.py, the wrapper handling in recompute.py, resolve_cp_group, nvtx_range.
• New here (not on dev): two-stage KDA gates, kda_disable_fp8, the KPool indexer, NoPE for the indexer / absorbed MLA, mla_disable_attention_fp8, the mHC precision flags, the mean output contract, full recompute with mHC, 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 DSA indexer weights projection precision controls #6842 (main: DSA indexer weights_proj precision controls), 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 layout-conversion package is ported, but the cp_partition_mode plumbing on TransformerConfig / PackedSeqParams that it depends on is not on main yet, so KDA is validated with CP = 1 here. CP > 1 lands together with the dev sync.
• KDA has no recurrent inference path (NotImplementedError), and gdn_pre_gated_delta_rule_fusion is not implemented for KDA.
• No functional test / golden values yet (see pre-checks). The change touches numerics, so Run functional tests is the right CI label.

Testing

Unit tests added or updated (GPU):

• 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/transformer/test_mhc_precision.py (new): mappings, aggregation and mixing against a pure-torch reference in both precision modes; fused mHC rejects the new flags.
• tests/unit_tests/models/test_hybrid_mhc.py: full recompute (uniform / block × fp32 / bf16 + FP32 mixing / bf16 fused) matches forward and backward without recompute; a wrapped residual layer does not double-count its input.
• tests/unit_tests/transformer/experimental_attention_variant/test_kpool_causal_tail.py (new): every query keeps its full causal prefix under the pool budget (packed / unpacked, strided queries, FP8 on / off); the FP8 input emulation matches a Hadamard-matrix reference.
• tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py: KPool projection dtypes and backward.
• tests/unit_tests/transformer/test_hyper_connection_recompute.py: drop the obsolete "mHC + full recompute is rejected" case.

pytest tests/unit_tests/ssm/test_kda_gate_precision.py \
       tests/unit_tests/transformer/test_mhc_precision.py \
       tests/unit_tests/models/test_hybrid_mhc.py \
       tests/unit_tests/transformer/experimental_attention_variant/test_kpool_causal_tail.py \
       tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py

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