[Refactor] Pass GLM DSA top-k IDs explicitly

[jayhenry]

Stack (bottom to top):

1. [Refactor] Switch gradient checkpointing to the non-reentrant implementation (1/3) #1979 feat/sac-nonlegacy-base → main
2. [Feature] Add the selective checkpointing engine (2/3) #1980 feat/sac-engine → feat/sac-nonlegacy-base
3. [Feature] Add recompute_cfg and per-model recompute declarations (3/3) #1981 feat/sac-recompute-cfg-v2 → feat/sac-engine
4. [Refactor] Pass GLM DSA top-k IDs explicitly #1988 feat/unify-offload → feat/sac-recompute-cfg-v2 ← you are here
5. [Feature] Keep DSA top-k out of checkpoint replay #1989 feat/dsa-topk-selective-checkpoint → feat/unify-offload

Base is PR3's branch. Review only this PR's own diff; the selective DSA indexer follow-up is isolated in PR5.

---

Summary

This is the stack-compatible port of #1978. The explicit DSA top-k dataflow and offload design stay the same, while the integration is adapted to the non-reentrant and selective checkpointing APIs introduced by #1979–#1981.

• Pass contiguous torch.int32 dsa_topk_ids explicitly through GLM decoder layers, intra-layer micro-batches, and MTP, removing the mutable SequenceContext cache and manual cross-layer lifecycle runtime.
• Keep all DSA-specific contracts under xtuner/v1/model/moe/glm52/: GLM52AttnOutputs, GLM dense/MoE decoder layers, DSA MLA, and GLM MTP layers/blocks. Generic attention, decoder, and MTP APIs no longer expose DSA fields or parameters.
• Add a model-agnostic _call_decoder_layer extension point to the main MoE and MTP loops. GLM-5.2 overrides it to arrange DSA inputs/results and contain the saved-tensor offload window without duplicating the full decoder stack.
• Use one saved-tensor offload window per main-stack decoder call for hidden states and the final shared-layer top-k IDs, matching tensors by storage so views share the same offload decision. MTP top-k offload remains excluded.
• Keep the sharing topology stateless: each main-stack source layer starts a new ID chain, shared layers consume the explicit tensor, and shared MTP logical depths forward the returned IDs directly.

Checkpoint-stack adaptation

• Preserve the keyed TypedDict decoder contract, list-based micro-batch API, and keyword checkpoint inputs instead of restoring the old reentrant-compatible flattened tuples.
• Extend only the GLM-derived attention and decoder outputs with required dsa_topk_ids; common outputs remain model-agnostic.
• Preserve the current attention/MLP/MoE checkpoint markers in the derived GLM paths, so non-reentrant full recompute and selective checkpoint regions keep the same boundaries.
• Sparse MLA backends consume the same contiguous int32 ID format across PyTorch, TileLang, cuDNN, and sequence-parallel paths.

Validation

• Full pre-commit suite passed: merge-conflict check, codespell, docformatter, pyupgrade, Ruff, Ruff format, mypy, and pydantic-extra-check.
• 28 focused tests passed, covering generic MoE behavior, GLM config/HF conversion, explicit DSA forward/backward, selective checkpoint replay, MTP compile/top-k offload, EP2 micro-batches, and the 8-GPU SP2 + EP4 + micro2 + compile + non-reentrant recompute + dual-offload path.
• A wheel build confirmed all six files in the new xtuner.v1.model.moe.glm52 package are included.

Performance evidence from #1978

These measurements are from the original isolated implementation in #1978; this stacked port reran the functional coverage above.

• 16K, 8×GPU AdamW: rank-0 effective TGS 10,893.7 vs 10,957.3 baseline (-0.58%; below the 11K absolute target); peak allocated memory 104.69 vs 105.26 GB.
• 32K, 8×GPU Muon SP2+EP2+micro2: full-rank sequence TGS 1,769.0 vs 1,599.3 baseline (+10.61%); peak allocated memory 101.87 vs 102.29 GB; VmPin 42.87 GiB in both runs.

Follow-up

PR #1989 keeps DSA top-k selection resident with selective checkpointing so backward replay does not recompute it.

Result

GLM-5.2 retains explicit, offloadable DSA top-k dataflow while the common attention, decoder, MTP, and checkpointing paths stay free of model-specific DSA APIs.

[HAOCHENYE]

Rebased this branch onto the current #1981 tip (`9d500b44`) — the SAC stack below was rewritten, so this branch had lost its base. New head is `9ba1c5f3`; the old one is preserved as `9d80c694` if you want to diff.

What changed in your commits: only the deleted checkpoint_record calls. The marker mechanism they belonged to was replaced by RecomputeTargetMap / KeptOps / KeptCallables, so those calls no longer resolve. Everything else — including the dsa_topk_ids threading, which is the substance of [Refactor] Pass GLM DSA top-k IDs explicitly — is byte-identical to what you wrote. Verifiable:

diff <(git diff 0e87277a 9d80c694) <(git diff 9d500b44 9ba1c5f3)

is exactly 6 lines, all of them checkpoint_record(...) or its import.

Worth knowing if you rebase this again: gh stack rebase does not work here. It computes the range from the merge-base with main, so it replays ~20 obsolete copies of the SAC commits and conflicts in 9 files. Use the explicit form instead:

git rebase --onto <new #1981 tip> <old base you branched from>

That replays only your own commits (4 conflicting files, all of them the checkpoint_record removal).

Checked: pre-commit passes on this PR's diff, both commits import cleanly, and the SAC unit tests pass. I did not run the GLM-5.2 engine tests — no checkpoint available here — so please confirm the DSA paths still behave.