From d2542776a0bff4092f2698413fce8f63e0d104c0 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Fri, 30 Oct 2026 10:22:35 +0000
Subject: [PATCH] fix(speculative): weigh DFlash2's selector loss like the
 backbone loss

The selector cross-entropy now uses the backbone loss's own per-position
weights (decay or D-PACE) and divides by the backbone's weight sum, misses
included. Before, every covered position counted equally and the term was
divided by the covered positions alone.

A block position's choice only matters if every earlier position was
accepted, which is what the backbone's weighting models, so the selector now
spends its capacity the same way. Dividing by the covered positions alone
also scaled the term up whenever the backbone's top-k recall dropped. This
matches the SGLang/SpecForge DFlash2 reference, and it is the objective the
released Muse-Spark-Open DFlash2 drafter was trained with.

The selector reads the backbone's mask and weights from _DFlashLossTerms
instead of rebuilding the mask, so the two cannot drift apart. The module
docstring still described the gold-token substitution that the strict top-k
replaced; it now describes the current objective.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../torch/speculative/plugins/hf_dflash2.py   | 58 +++++++++----------
 .../speculative/plugins/test_hf_dflash2.py    | 32 +++++++++-
 2 files changed, 58 insertions(+), 32 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py
index 6060923b546..b0a14d16937 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash2.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash2.py
@@ -29,9 +29,11 @@
 - Take the backbone's top-k candidates per block position.
 - Score each candidate against its *teacher-forced* predecessor token, so the
   positions train in parallel exactly as the backbone does.
-- When the gold token is missing from the top-k, substitute it into the last
-  candidate slot. Without this the selector sees no positive class on the hard
-  positions and never learns those edges.
+- Supervise only the strict top-k, the candidate set serving builds. A gold token
+  outside it is a backbone recall miss, not a selector example, and carries no
+  selector gradient.
+- Weight each position as the backbone loss does (decay or D-PACE) and divide by the
+  backbone's weight sum, misses included.
 """
 
 import torch
@@ -106,7 +108,9 @@ def get_exporter(self):
 
         return DFlash2Exporter(self)
 
-    def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mask):
+    def _selector_loss(
+        self, logits, target_ids, hidden, predecessor_ids, supervised_mask, loss_weights, weight_sum
+    ):
         """Cross-entropy over the selector's candidate set, and its top-1 accuracy.
 
         Args:
@@ -114,7 +118,10 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas
             target_ids: Gold token ids ``[B, N, block_size]``.
             hidden: Backbone hidden states ``[B, N, block_size, H]``.
             predecessor_ids: Teacher-forced predecessor ids ``[B, N, block_size]``.
-            weight_mask: Per-position loss weights ``[B, N, block_size]``.
+            supervised_mask: The positions the backbone loss supervises, unweighted
+                ``[B, N, block_size]``.
+            loss_weights: The backbone loss's per-position weights ``[B, N, block_size]``.
+            weight_sum: The backbone loss's denominator, the sum of ``loss_weights``.
 
         Returns:
             ``(loss, accuracy, coverage)`` — coverage is the fraction of supervised
@@ -130,7 +137,7 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas
         # token the backbone did not propose is a backbone recall failure, not a
         # selector classification example: substituting it in would teach the selector
         # to override the unary ranking on a set it will never be shown. Those
-        # positions carry no selector gradient and leave the denominator instead.
+        # positions carry no selector gradient.
         gold_matches = candidate_ids == target_ids.unsqueeze(-1)
         gold_in_topk = gold_matches.any(dim=-1)
         gold_slot = gold_matches.long().argmax(dim=-1)
@@ -139,26 +146,27 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas
             candidate_ids, unary_logits, hidden, predecessor_ids
         )
 
-        covered = weight_mask * gold_in_topk.to(weight_mask.dtype)
-        flat_weights = covered.reshape(-1)
-        denominator = flat_weights.sum() + 1e-6
+        in_topk = gold_in_topk.to(loss_weights.dtype)
         per_token = F.cross_entropy(
             selector_logits.float().reshape(-1, top_k),
             gold_slot.reshape(-1),
             reduction="none",
         )
-        loss = (per_token * flat_weights).sum() / denominator
+        # The backbone's denominator, misses included: dividing by the covered positions
+        # alone would scale the term up whenever backbone recall drops.
+        loss = (per_token * (loss_weights * in_topk).reshape(-1)).sum() / weight_sum
 
         with torch.no_grad():
+            covered = (supervised_mask * in_topk).reshape(-1)
             chosen = selector_logits.argmax(dim=-1).reshape(-1)
-            accuracy = (
-                (chosen == gold_slot.reshape(-1)).float() * flat_weights
-            ).sum() / denominator
+            accuracy = ((chosen == gold_slot.reshape(-1)).float() * covered).sum() / (
+                covered.sum() + 1e-6
+            )
             # Coverage keeps the full supervised mask as its denominator: it measures
             # how often the selector was given a solvable problem at all.
             # clamp, not an epsilon: a fully covered batch must read exactly 1.0.
-            supervised = weight_mask.reshape(-1).sum()
-            coverage = flat_weights.sum() / supervised.clamp(min=1.0)
+            supervised = supervised_mask.reshape(-1).sum()
+            coverage = covered.sum() / supervised.clamp(min=1.0)
         # Detached tensors, not Python scalars: .item() would force a CPU-GPU sync on
         # every training step. The trainer converts them at the logging boundary.
         return loss, accuracy.detach(), coverage.detach()
@@ -238,24 +246,12 @@ def _compute_loss(
 
         offsets = torch.arange(block_size, device=device).view(1, 1, -1)
         label_indices = anchor_positions.unsqueeze(-1) + offsets
-        valid_label = label_indices < seq_len
         safe_label_indices = label_indices.clamp(max=seq_len - 1)
         expanded_ids = input_ids.unsqueeze(1).expand(-1, n_blocks, -1)
         target_ids = torch.gather(expanded_ids, 2, safe_label_indices)
 
-        # Same supervision mask as the backbone loss: valid block, in bounds, not the
-        # anchor slot, and inside the answer span. Position weighting (decay/D-PACE) is
-        # deliberately not applied — it shapes *where* the backbone spends capacity,
-        # while the selector should learn every position's transition equally.
-        weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float()
-        weight_mask = weight_mask * valid_label.float()
-        weight_mask = weight_mask * (offsets > 0).float()
-        weight_mask = weight_mask * torch.gather(
-            loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
-        )
-
         # Teacher-forced predecessor of block position k is the real token at anchor+k-1.
-        # Only offsets 1..block_size-1 are supervised (weight_mask zeroes slot 0), so the
+        # Only offsets 1..block_size-1 are supervised (the backbone's mask zeroes slot 0), so the
         # first supervised position, k=1, has the anchor's own token as its predecessor.
         # Slot 0's entry here resolves to anchor-1 and never contributes.
         #
@@ -264,12 +260,16 @@ def _compute_loss(
         # that hands greedy_path the full 0..block_size-1 candidate set is off by one.
         predecessor_ids = torch.gather(expanded_ids, 2, (safe_label_indices - 1).clamp(min=0))
 
+        # The backbone's own mask and position weights: a block position's choice only
+        # matters if every earlier one was accepted, which decay / D-PACE model.
         selector_loss, selector_accuracy, selector_coverage = self._selector_loss(
             logits.reshape(bsz, n_blocks, block_size, -1),
             target_ids,
             draft_hidden.reshape(bsz, n_blocks, block_size, -1),
             predecessor_ids,
-            weight_mask,
+            terms.supervised_mask.view_as(target_ids),
+            terms.weights.view_as(target_ids),
+            terms.weight_sum,
         )
         self._selector_metrics = {
             "selector_accuracy": selector_accuracy,
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
index c43677818c0..e1ea8c10dc0 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
@@ -397,9 +397,11 @@ def _selector(self, vocab=16, rank=4, top_k=3, hidden=8):
             torch.nn.init.normal_(selector.successor_codebook, std=0.5)
         return selector
 
-    def _selector_terms(self, selector, target_ids, candidate_ids):
+    def _selector_terms(self, selector, target_ids, candidate_ids, loss_weights=None):
         """Drive HFDFlash2Model._selector_loss with a hand-built candidate set."""
         b, length, k = candidate_ids.shape
+        supervised = torch.ones(b, length, dtype=torch.double)
+        weights = supervised if loss_weights is None else loss_weights
         model = HFDFlash2Model.__new__(HFDFlash2Model)
         model.dflash_module = type("_M", (), {"candidate_selector": selector})()
         logits = torch.zeros(b, length, 32, dtype=torch.double)
@@ -410,9 +412,13 @@ def _selector_terms(self, selector, target_ids, candidate_ids):
             model,
             logits,
             target_ids,
-            torch.randn(b, length, 8, dtype=torch.double),
+            torch.randn(
+                b, length, 8, generator=torch.Generator().manual_seed(1), dtype=torch.double
+            ),
             torch.zeros(b, length, dtype=torch.long),
-            torch.ones(b, length, dtype=torch.double),
+            supervised,
+            weights,
+            weights.sum(),
         )
 
     def test_a_miss_carries_no_selector_gradient(self):
@@ -453,6 +459,26 @@ def test_a_covered_position_does_train_the_selector(self):
         loss.backward()
         assert selector.successor_codebook.grad.abs().sum() > 0.0
 
+    def test_positions_weigh_like_the_backbone_and_misses_stay_in_the_denominator(self):
+        """The term is the backbone-weighted mean over every supervised position.
+
+        A miss adds nothing to the numerator but keeps its weight in the denominator, so
+        the term does not grow as backbone recall drops.
+        """
+        selector = self._selector(vocab=32, top_k=2)
+        candidates = torch.tensor([[[5, 6], [5, 6], [5, 6]]])
+        targets = torch.tensor([[5, 6, 9]])  # two covered positions, one miss
+        weights = torch.tensor([[1.0, 3.0, 2.0]], dtype=torch.double)
+
+        loss, _, _ = self._selector_terms(selector, targets, candidates, weights)
+        per_position = [
+            self._selector_terms(selector, targets, candidates, torch.eye(3)[i : i + 1].double())[0]
+            for i in range(3)
+        ]
+        assert float(per_position[2]) == 0.0
+        expected = sum(w * term for w, term in zip(weights[0], per_position)) / weights.sum()
+        torch.testing.assert_close(loss, expected)
+
     def test_greedy_path_position_zero_is_seeded_by_the_anchor(self):
         """``greedy_path`` position 0 scores against the anchor, i.e. block offset 1.