[h-guo18]

Stack (merge in this order; each PR's base is the branch of the one before it). [7/7] reproduces nvidia/Muse-Spark-Open-816B-A42B-DFlash2 and needs all of them.

1. [1/4] fix(speculative): four training fixes behind the Muse-Spark-Open DFlash2 recipe #2617: draft RoPE frequencies stay in fp32
2. [2/7] fix(speculative): make streaming hidden-state capture work on vLLM 0.29 #2522: streaming hidden-state capture on vLLM >= 0.29
3. [3/7] fix(speculative): keep the streaming corpus memory-mapped instead of copying it per rank #3142: streaming corpus stays memory-mapped
4. [4/7] fix(speculative): weigh DFlash2's selector loss like the backbone loss #3143: DFlash2 selector loss weighted like the backbone loss <- this PR
5. [2/4] perf(speculative): deliver training batches out of order; compile + FlexAttention in the DFlash2 recipe #3001: out-of-order training batches; compile + FlexAttention in the DFlash2 recipe
6. [3/4] perf(speculative): fuse the DFlash-family loss with the lm_head projection #3010: fused DFlash-family loss
7. [4/4] feat(speculative): Muse-Spark-Open DFlash2 recipes that reproduce the released drafter #3144: Muse-Spark-Open DFlash2 recipes

---

What does this PR do?

Type of change: Bug fix

DFlash2's selector cross-entropy now weighs block positions the way the backbone loss does, and shares the backbone's denominator:

• Weights. Each covered position's selector cross-entropy is multiplied by the backbone's per-position loss weight (decay or D-PACE), read from _DFlashLossTerms.weights. Before, every covered position counted equally.
• Denominator. The sum is divided by the backbone loss's weight sum, _DFlashLossTerms.weight_sum, which also counts the positions whose gold token is outside the top-k. Before, it was divided by the covered positions alone, so the term grew whenever the backbone's top-k recall dropped.
• The selector now reads the backbone's supervision mask from _DFlashLossTerms.supervised_mask instead of rebuilding it, so the two cannot drift apart.

Why. A block position's choice only matters if every earlier position was accepted, and the backbone's weighting models exactly that. With this change the selector spends its capacity where acceptance is decided. This is how the SGLang/SpecForge DFlash2 reference trains its selector, and it is the objective that the released nvidia/Muse-Spark-Open-816B-A42B-DFlash2 drafter was trained with.

What I have not measured. I have no A/B of this change on its own. main already trains the selector on the strict top-k (no gold-token substitution) and zero-initializes its successor codebook. Together with this change, that is the objective that beat the old gold-substituting one by +7.8% (greedy) and +10.1% (temperature 1) MT-Bench acceptance length, in from-scratch Muse-Spark-Open runs of 550K samples each. How that gain splits between the strict top-k and the weighting here, I did not measure.

The module docstring still described the gold-token substitution that the strict top-k replaced. It now describes the current objective.

No CHANGELOG entry: DFlash2 first ships in 0.48.0, which is still at a release candidate.

Usage

No API or config change. The selector term's value changes, so a run's logged loss will not match a run on the previous objective.

Testing

• New unit test TestCandidateSelectorAlignment::test_positions_weigh_like_the_backbone_and_misses_stay_in_the_denominator in tests/unit/torch/speculative/plugins/test_hf_dflash2.py. It checks that the term is the backbone-weighted mean over every supervised position, a miss included. With the previous objective restored, it fails and the other 33 tests in the file pass.
• pytest of the DFlash-family unit tests (test_hf_dflash.py, test_hf_dflash2.py, test_fakebase.py, test_dflash_fused_head.py, test_hf_lilicorr.py, test_modeling_final_norm.py) on the top of the stack: passed with transformers 5.5.4 and 5.18.0 (CPU).
• End to end, see [7/7], which trains with this objective.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ (no API change; the selector term's value changes)
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A (DFlash2 has not shipped yet)
• Did you get Claude approval on this PR?: ❌

Additional Information

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[h-guo18]

Folded into #2617 together with #2617's own fix and the other two small fixes, rebased onto main after #2757. It is now [1/4] of the Muse-Spark-Open DFlash2 recipe stack. The commits are the same, re-signed, with the selector one ported onto #2757's helpers; see #2617 for the combined description and testing.