feat: DSA conversion (dense-MLA -> DeepSeek Sparse Attention), with a Kimi K2 example

[faresobeid]

Stacked on #3065 (Kimi K2 support) — base branch is feat/kimi-k2-support, not main. Diff shown here is only this PR's changes on top of that one.

Summary

Machinery to convert an existing dense-MLA checkpoint to DeepSeek Sparse Attention (DSA) via continued pretraining, mirroring the two-stage recipe DeepSeek (arXiv:2512.02556), GLM (arXiv:2602.15763), and IndexCache (arXiv:2603.12201) used, plus recovery afterward with the existing opd algorithm.

• layers/dsa.py (new): SparseMlaAttention/Indexer relocated out of glm_moe_dsa (now a thin consumer) so any dense-MLA family can build a DSA sibling on it. use_sparse_attn toggles dense-causal vs. real DSA top-k attention over the same MLA weights — that's what lets one checkpoint move through both conversion stages without a format migration in between.
• Indexer training: the existing compute_sparse_indices stays exactly as-is (non-differentiable — top-k has no gradient regardless of what feeds it). A new differentiable Indexer.score() + compute_indexer_kl_loss train it via KL against the real attention distribution (dense stage) or the actual sparse-selected keys (sparse stage).
• Bootstrap: a dense checkpoint has no indexer.* keys, and prime-rl's checkpoint loader is strict (Missing key in checkpoint state_dict). strip_indexer_from_state_dict + Indexer._init_indexer_parameters, wired into load_dcp_from_hf behind model.debug.random_init_indexer, mirror the exact strip/reinit pattern already used for LoRA adapters.
• freeze_all_except_indexer (+ model.freeze_all_except_indexer): the paper recipe's indexer warm-up stage needs everything but the indexer frozen — existing freeze_sparse_indexer only ever freezes the indexer itself, the opposite of what this stage needs.
• RawTextDataConfig/RawTextDataset ([data] type = "raw_text"): a continued-pretraining data path for teams reproducing DeepSeek-V3.2's/GLM-5's own from-scratch conversion recipe (raw text, not chat data). Not required in general — see below.
• kimi_k2_dsa: new family pairing with feat(models): add Kimi K2 + K2.7 (text-only) support #3065's kimi_k2, built entirely on layers/dsa.py (~30 lines of config, no new attention code) — proof the shared module actually generalizes beyond the family it shipped with.
• docs/advanced.md: minimal "DSA Conversion" section with a worked Kimi K2 example (bootstrap → indexer warm-up → sparse adaptation → opd recovery).

Correction from an earlier revision: neither the indexer-KL loss nor the dense-mode forward pass consults loss_mask, so ordinary SFTDataset already works for the whole conversion recipe unmodified — confirmed by re-reading arXiv:2603.12201 (IndexCache) directly, whose own from-scratch conversion of GLM-4.7-Flash used plain SFT data for both stages (1,000 warm-up + 4,000 sparse-training steps, 200K context), not a raw-text corpus. The docs and worked example now lead with SFT data as the practical default; RawTextDataConfig remains for reproducing DeepSeek/GLM-5's specific recipe.

Verification

• GPU: dense-mode output matches the existing sparse-attention path exactly when top-k covers the full sequence, for both glm_moe_dsa and kimi_k2_dsa — the second confirms layers/dsa.py generalizes rather than just working for the family it was extracted from.
• GPU: the indexer-KL loss measurably trains a randomly-initialized indexer (loss decreases over 20 steps against a fixed target).
• CPU: bootstrap strip/reinit and freeze_all_except_indexer are unit-tested directly.
• Full trainer test suite: no new failures (same pre-existing, unrelated failures reproduce on a clean checkout of this base branch — a flash-attn/fp32 issue affecting several existing model families, unrelated to this change).
• Not done: no real Kimi K2 (or GLM-5) checkpoint was available to run the actual conversion recipe end-to-end — this validates the mechanism, not a specific trained result.

🤖 Generated with Claude Code