[Refactor] Unify SAC checkpointing and GLM DSA dataflow

[jayhenry]

Summary

• Unify activation checkpointing around one PyTree-aware reentrant boundary so structured micro-batch inputs and outputs stay on the real autograd and offload path.
• Pass GLM DSA top-k IDs explicitly through decoder and MTP outputs, removing the mutable SequenceContext cache lifecycle.
• Reuse frozen source-indexer results from checkpoint-local FIFO frames during replay.
• Keep FSDP wrappers outside recompute and dense checkpoint boundaries outside compiled graphs.

Regression coverage

• compile + top-k offload + shared-weight MTP with depth > 1; validates pinned-memory offload and restore, and that the source indexer is not recomputed
• EP > 1 + intra-layer micro-batch 2
• FSDP checkpointing, torch.compile, FP8, activation offload, and nested PyTree inputs and outputs

Validation

• Post-rebase on latest upstream/main: 16 passed
  • tests/model/test_recompute.py
  • tests/model/test_glm52_mtp_checkpoint_repro.py
  • tests/model/test_fsdp_checkpoint.py
• Qwen3.5 full 40-layer stack5 regression: step losses and grad norms matched upstream; steady sequence TGS +0.36%; peak allocated memory -0.14 GB.
• GLM AdamW production-shaped and Muon EP2/micro2 training regressions completed.