[HollowMan6]

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

What does this PR do?

Add KPool (pooled-key) DSA indexer support, NoPE support for the DSA indexer and absorbed MLA, and opt-in precision controls for matching serving-time precision under FP8 training.

This PR uses the following configuration flags:

• dsa_indexer_qk_proj_disable_quantization
• dsa_indexer_kpool_use_quantization
• mla_proj_disable_quantization

The DSA weights projection remains on its generic path. Its precision is controlled separately by the weights-projection options introduced by #6842; this PR does not force it to FP32 or unconditionally disable its quantization.

Changes

• dsa_indexer_kpool > 1: compress keys in groups of kpool tokens using a BF16 per-token compression gate and an FP32 per-slot positional bias. Pool selection is performed over compressed keys and expanded back to token indices. dsa_indexer_kpool_always_select_tail appends each query's incomplete causal pool. Packed THD inputs restart pooling at document boundaries. The fused cuDNN DSA path is bypassed for KPool.

• dsa_indexer_kpool_use_quantization: optionally emulate the serving indexer's quantized KPool scoring inputs. When disabled, KPool scoring uses the unquantized configured dtype.

• dsa_indexer_qk_proj_disable_quantization: optionally keeps the DSA indexer's query/key projection initialization and GEMMs unquantized, typically in BF16. This flag does not control the weights projection.

• NoPE support: qk_pos_emb_head_dim == 0 is supported by DSAIndexer and AbsorbedMLA without constructing or applying rotary embeddings.

• mla_proj_disable_quantization: optionally keeps absorbed-MLA projections (q_down, q_up, kv_down, and output projection) unquantized. Core-attention precision is unchanged.

Motivation

The actor and rollout implementations must use matching precision for modules whose outputs affect discrete DSA top-k selection or MLA attention scores. These flags allow model bridges to explicitly reproduce the serving precision split while leaving unrelated projections and the rest of the model on their normal quantization path.

Relationship to other PRs

• Related: Add DSA indexer weights projection precision controls #6842 for DSA weights-projection precision controls.
• Related: [dev] feat(moe): Add KDA Low-Rank Projections and Extend MLA Support #6877 for MLA NoPE support.
• Related: GLM-5.2 training support tracking #6392 for GLM-5.2/GLM-5.3 training support tracking.
• Sibling PRs: Add mHC precision controls, mean output contract, and full activation recompute with mHC #7521 and Guard blockwise FP8 backward (E5M2) wgrad against all-zero gradient blocks #7523.

Testing

Relevant tests cover:

• KPool causal-tail selection and packed/unpacked inputs.
• DSA projection precision, backward behavior, FP32 k-norm, and KPool activation cleanup.
• Multi-TP DSA consistency.
• Absorbed MLA NoPE behavior.
• Kernel determinism coverage for modified fused paths.

Example:

python -m torch.distributed.run --standalone --nproc-per-node 8 \
  -m pytest -q \
  tests/unit_tests/transformer/experimental_attention_variant/test_kpool_causal_tail.py \
  tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py

python tools/check_kernel_determinism_coverage.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 #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.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[buptzyb]

I think we have to make sure our changes don't break existing models' behaviors. Thanks!

[guihong-nv]

/ok to test 073bf15

[HollowMan6]

@guihong-nv Fixed the deterministic lint, would you mind retriggering the CI? Thanks!

[Phlip79]

/ok to test c6deff8

[guihong-nv]

/ok to test 0e92955

[nemo-automation-bot]

🔄 Merge queue validation started!

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