Feature Request: DSA/DSv4 Indexer Replay for RL
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Summary
Add replay support for DSA / DeepSeek-V4 hybrid-attention indexer top-k decisions during RL training.
Megatron already has MoE router replay support for R3-style RL. We should add something similar for sparse-attention indexer decisions in DSA / DSv4 hybrid attention.
## Motivation
For RL, rollout inference and train-side recomputation may run through different kernels, precisions, batching layouts, CUDA graph paths, or packed-sequence layouts. In DSA/DSv4 sparse attention, the indexer makes a discrete top-k selection of which KV positions each query token attends to.
Small numerical differences can cause rollout inference and training to choose different indexer top-k positions for the same tokens. That means train-side logprobs may be computed through a different sparse-attention graph than the one used during rollout, creating train/rollout mismatch.
This is the same instability that MoE R3 addresses for expert routing, but applied to sparse-attention indexer selections.
## Proposed Solution
Add an opt-in replay path, e.g. `dsa_enable_indexer_replay`, that supports record mode, replay mode, and RL frameworks to carry recorded indexer metadata with rollout data.
This should also be tested with downstream RL frameworks like verl.
## Related Work
MoE Router Replay / R3 issue: https://github.com/NVIDIA/Megatron-LM/issues/4168
MCore RouterReplay PR: https://github.com/NVIDIA/Megatron-LM/pull/2693
DSA support: https://github.com/NVIDIA/Megatron-LM/pull/2440
Contributor guide
Assessment
This issue has not been assessed yet.