[Feature] Support GLM-5.2 source-layer DSA indexer training

[sjyouuuuug]

Motivation

Enable training of GLM-5.2 source-layer DSA indexers with attention-derived KL supervision.

Changes

• Add a cuDNN DSA indexer KL-loss operator and custom backward integration.
• Separate differentiable indexer feature projection from no-grad Top-k selection.
• Construct sparse attention teacher and indexer student distributions.
• Aggregate source-layer indexer losses at the model boundary.
• Keep physical MTP indexers frozen.
• Add a GLM-5.2 training configuration and launch script.
• Add tests for KL construction, gradient boundaries, score scaling, query masking, source/shared-layer behavior, and MTP freezing.

Validation

• Indexer loss CPU tests: 10 passed.
• Source-indexer attention tests: 10 passed.
• Multi-GPU tiny training completed all configured training steps.
• Indexer loss decreased during the 300-step training run.

Current limitations

• IndexShare reuses source-layer Top-k indices, but shared-layer teacher supervision is not included.
• Indexer training currently requires:
  • SP_SIZE=1
  • INTRA_LAYER_MICRO_BATCH=1
  • the cudnn_dsa backend

Results

[jayhenry]

两个weights能否统一？

[sjyouuuuug]

已经统一了，现在 weights_proj 只计算一次，并在 DSAIndexerFeatures 中保存一份 raw weights，不再维护 selection_weights 和 training_weights 两份张量。

[jayhenry]

对于 share layer 则不训练 indexer？

[sjyouuuuug]

现在只训练 source layer 的 indexer，后面只复用，不算 kl

[jayhenry]

是否有参考来源？

[sjyouuuuug]

主要参考 https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/experimental_attention_variant/dsa.py 中的compute_dsa_indexer_loss，算法的逻辑上是比较近似的，主 attention 的 query 计算 teacher attention distribution，通过 KL(teacher || indexer) 训练 indexer。sparse 模式下使用 Top-k mask 限制 token。