From d9658d5b0eefbbf97751c79f937b1f3271fb330e Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 07:39:59 +0000
Subject: [PATCH 01/15] refactor(speculative): read the teacher by label
 position; one final-norm core

The next-token shift -- the base logits that predict the token at position i sit at i - 1 --
was spelled out at each of _teacher_logits' three callers: DFlash KD and DSpark as
(safe_label_indices - 1).clamp(min=0), LiLiCorr as anchor + s for the token at anchor + 1 + s.
A wrong shift trains without an error, just worse, so it now lives once: _teacher_logits
takes the label positions and applies the shift itself, and the callers pass what they
supervise. DFlash2's predecessor gather uses the same offset for a different quantity (the
previous token, not the teacher) and keeps its own.

_maybe_apply_base_final_norm took the producer's whole outputs dict, so _teacher_logits
wrapped its one flag in a single-key dict to call it. The bool now has its own function,
_apply_base_final_norm, which the dict reader delegates to; EAGLE keeps calling the dict
reader, unchanged.

Behavior is unchanged: against this branch's parent, offline DFlash KD, LiLiCorr (penalty and
calibration) and DSpark give bitwise-equal losses and gradients in fp32 and bf16, with pre-
and post-norm captures stored in float64. The teacher tests now pass labels and expect row
label - 1, and fail without the shift; a new test pins that label 0 reads row 0 rather than
wrapping to the last row.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../torch/speculative/plugins/hf_dflash.py    | 22 ++++-----
 .../torch/speculative/plugins/hf_dspark.py    |  6 +--
 .../torch/speculative/plugins/hf_lilicorr.py  | 11 ++---
 .../plugins/modeling_final_norm.py            | 29 +++++++----
 .../speculative/plugins/test_hf_dflash.py     | 48 +++++++++++--------
 5 files changed, 64 insertions(+), 52 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index c1c3389e428..7d642602768 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -97,7 +97,7 @@
     _LM_HEAD_PATHS,
     find_final_norm_path,
 )
-from .modeling_final_norm import _maybe_apply_base_final_norm
+from .modeling_final_norm import _apply_base_final_norm
 
 logger = logging.getLogger(__name__)
 
@@ -745,13 +745,15 @@ def _build_generate_swa_mask(self, ctx_len, bsz, dtype, device):
         return attn_mask
 
     @torch.no_grad()
-    def _teacher_logits(self, base_outputs, positions, token_ids=None):
-        """Base-model logits at ``positions`` ([B, ...] sequence indices) -> [B, ..., vocab].
+    def _teacher_logits(self, base_outputs, label_positions, token_ids=None):
+        """The base model's logits for the tokens at ``label_positions`` ([B, ...]) -> [B, ..., vocab].
 
-        Only the requested rows of the base hidden go through the final norm and lm_head, so
-        no full-sequence logits are built from it; ``token_ids`` ([B, ..., k]) narrows the
+        Next-token logits: the ones that predict the token at position ``i`` sit at ``i - 1``.
+        Only those rows of the base hidden go through the final norm and lm_head, so no
+        full-sequence logits are built from it; ``token_ids`` ([B, ..., k]) narrows the
         projection to just those vocab entries.
         """
+        positions = (label_positions - 1).clamp(min=0)
         if base_outputs.logits is not None:
             if token_ids is None:
                 return _gather_rows(base_outputs.logits, positions)
@@ -767,11 +769,7 @@ def _teacher_logits(self, base_outputs, positions, token_ids=None):
         # A producer can store the hidden states in a wider dtype than the target's weights.
         # Cast before the final norm too, which online training runs in the target's dtype.
         rows = _gather_rows(base_outputs.base_hidden, positions).to(lm_head.weight.dtype)
-        rows = _maybe_apply_base_final_norm(
-            rows,
-            {"base_hidden_prenorm": base_outputs.base_hidden_prenorm},
-            self._base_model_norm,
-        )
+        rows = _apply_base_final_norm(rows, base_outputs.base_hidden_prenorm, self._base_model_norm)
         if token_ids is None:
             return lm_head(rows)
         logits = torch.einsum("...h,...kh->...k", rows, lm_head.weight[token_ids])
@@ -882,9 +880,7 @@ def _compute_loss(
 
         if valid_count > 1.0:
             if kd:
-                # KD loss: teacher logits for token anchor+k are at position anchor+k-1
-                teacher_indices = (safe_label_indices - 1).clamp(min=0)
-                flat_teacher = self._teacher_logits(base_outputs, teacher_indices)
+                flat_teacher = self._teacher_logits(base_outputs, safe_label_indices)
                 flat_teacher = flat_teacher.reshape(flat_logits.shape)
                 target_soft = torch.softmax(flat_teacher, dim=-1)
                 draft_logsoft = torch.log_softmax(flat_logits, dim=-1)
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index 7c81f4cc478..58ae972d628 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -244,10 +244,8 @@ def _compute_dspark_loss(
             metrics = {"ce_loss": 0.0, "l1_loss": 0.0, "confidence_loss": 0.0, "base_accuracy": 0.0}
             return loss, 0.0, metrics
 
-        # Aligned target distribution: base-model logits that predict token anchor+k+1
-        # sit at position anchor+k (= label index - 1).
-        teacher_indices = (safe_label_indices - 1).clamp(min=0)
-        flat_teacher = self._teacher_logits(base_outputs, teacher_indices).reshape(-1, vocab)
+        # Aligned target distribution: the base model's own prediction of each label.
+        flat_teacher = self._teacher_logits(base_outputs, safe_label_indices).reshape(-1, vocab)
 
         # Term 1: cross-entropy on the corrected (final) logits.
         ce_per_token = F.cross_entropy(flat_final, flat_targets, reduction="none")
diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py
index 42f6a8bbed5..c77bbcd9f7d 100644
--- a/modelopt/torch/speculative/plugins/hf_lilicorr.py
+++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py
@@ -545,14 +545,13 @@ def _candidate_target_logits(
         batch_blocks, num_slots, topk = candidate_ids.shape
         bsz, n_blocks = anchor_positions.shape
         seq_len = base_outputs.target_hidden.shape[1]
-        # The target's next-token logits that predict the token at anchor+1+s sit at
-        # position anchor+s. Only the [blocks, slots, k] candidate logits are computed,
-        # never whole vocabulary rows.
-        slot_offsets = torch.arange(num_slots, device=candidate_ids.device)
-        positions = (anchor_positions.unsqueeze(-1) + slot_offsets).clamp(min=0, max=seq_len - 1)
+        # Slot s scores the token at anchor+1+s. Only the [blocks, slots, k] candidate logits
+        # are computed, never whole vocabulary rows.
+        slot_offsets = torch.arange(1, num_slots + 1, device=candidate_ids.device)
+        labels = (anchor_positions.unsqueeze(-1) + slot_offsets).clamp(max=seq_len - 1)
         candidate_target_logits = self._teacher_logits(
             base_outputs,
-            positions,
+            labels,
             token_ids=candidate_ids.reshape(bsz, n_blocks, num_slots, topk),
         )
         return candidate_target_logits.reshape(batch_blocks, num_slots, topk).float()
diff --git a/modelopt/torch/speculative/plugins/modeling_final_norm.py b/modelopt/torch/speculative/plugins/modeling_final_norm.py
index d35badb5aa8..28b1ff5a0db 100644
--- a/modelopt/torch/speculative/plugins/modeling_final_norm.py
+++ b/modelopt/torch/speculative/plugins/modeling_final_norm.py
@@ -123,21 +123,20 @@ def _select_final_norm_type(model_type: str | None, base_cfg=None) -> str | None
     return _FINAL_NORM_TYPE_BY_MODEL_TYPE.get(model_type or "")
 
 
-def _maybe_apply_base_final_norm(hidden, base_model_outputs, base_model_norm):
-    """Re-apply the base model's final norm to ``hidden`` before lm_head, per the producer.
+def _apply_base_final_norm(hidden, prenorm, base_model_norm):
+    """Re-apply the base model's final norm to ``hidden`` before lm_head, if it was captured before it.
 
-    When the offline/streaming producer captured a *pre-final-norm* hidden it sets
-    ``base_hidden_prenorm=True`` in ``base_model_outputs``; the consumer must re-apply the base
-    final norm to reconstruct correct logits. Shared by the DFlash and EAGLE offline forwards.
+    When the offline/streaming producer captured a *pre-final-norm* hidden (``prenorm``), the
+    consumer must re-apply the base final norm to reconstruct correct logits.
 
-    - ``base_hidden_prenorm`` falsy (post-norm capture, e.g. HF/TRT-LLM): return ``hidden`` as-is.
-    - ``base_hidden_prenorm`` true and a norm is available: return the normed hidden.
-    - ``base_hidden_prenorm`` true but ``base_model_norm is None``: raise — we cannot reconstruct
-      correct logits, and silently feeding an un-normed hidden into lm_head would corrupt the
+    - ``prenorm`` false (post-norm capture, e.g. HF/TRT-LLM): return ``hidden`` as-is.
+    - ``prenorm`` true and a norm is available: return the normed hidden.
+    - ``prenorm`` true but ``base_model_norm is None``: raise — we cannot reconstruct correct
+      logits, and silently feeding an un-normed hidden into lm_head would corrupt the
       distillation target. The base ``model_type`` is not enabled in
       ``_FINAL_NORM_TYPE_BY_MODEL_TYPE``.
     """
-    if not base_model_outputs.get("base_hidden_prenorm", False):
+    if not prenorm:
         return hidden
     if base_model_norm is None:
         raise RuntimeError(
@@ -149,3 +148,13 @@ def _maybe_apply_base_final_norm(hidden, base_model_outputs, base_model_norm):
             "matching norm class in _FINAL_NORM_CLASSES if it needs a new norm flavor."
         )
     return base_model_norm(hidden)
+
+
+def _maybe_apply_base_final_norm(hidden, base_model_outputs, base_model_norm):
+    """:func:`_apply_base_final_norm`, reading ``prenorm`` from the producer's outputs dict.
+
+    The producer sets ``base_hidden_prenorm=True`` in ``base_model_outputs`` for a pre-norm
+    capture; older dumps omit it and are post-norm.
+    """
+    prenorm = bool(base_model_outputs.get("base_hidden_prenorm", False))
+    return _apply_base_final_norm(hidden, prenorm, base_model_norm)
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index ff8931149e5..ae8337daba5 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -1446,7 +1446,10 @@ def test_trailing_padding_blocks_do_not_change_the_loss(self):
 
 
 class TestTeacherLogits:
-    """The base distribution is projected only at the rows, and vocab entries, asked for."""
+    """The base distribution is projected only at the rows, and vocab entries, asked for.
+
+    It is the next-token distribution, so label ``i`` reads row ``i - 1``.
+    """
 
     @staticmethod
     def _model():
@@ -1455,9 +1458,9 @@ def _model():
         return model
 
     @staticmethod
-    def _positions(bsz=2):
+    def _labels(bsz=2):
         torch.manual_seed(0)
-        return torch.randint(0, SEQ_LEN, (bsz, 3, BLOCK_SIZE))
+        return torch.randint(1, SEQ_LEN, (bsz, 3, BLOCK_SIZE))
 
     @staticmethod
     def _rows(full, positions):
@@ -1471,45 +1474,52 @@ def _hidden(model, bsz=2):
     def test_matches_the_full_sequence_projection(self):
         model = self._model()
         hidden = self._hidden(model)
-        positions = self._positions()
-        got = model._teacher_logits(DFlashBaseModelOutput(None, base_hidden=hidden), positions)
-        want = self._rows(model._base_model_lm_head(hidden), positions)
+        labels = self._labels()
+        got = model._teacher_logits(DFlashBaseModelOutput(None, base_hidden=hidden), labels)
+        want = self._rows(model._base_model_lm_head(hidden), labels - 1)
         torch.testing.assert_close(got, want)
 
     def test_prenorm_hidden_gets_the_base_final_norm(self):
         model = self._model()
         hidden = self._hidden(model)
-        positions = self._positions()
+        labels = self._labels()
         outputs = DFlashBaseModelOutput(None, base_hidden=hidden, base_hidden_prenorm=True)
-        want = self._rows(model._base_model_lm_head(model._base_model_norm(hidden)), positions)
-        torch.testing.assert_close(model._teacher_logits(outputs, positions), want)
+        want = self._rows(model._base_model_lm_head(model._base_model_norm(hidden)), labels - 1)
+        torch.testing.assert_close(model._teacher_logits(outputs, labels), want)
 
     def test_token_ids_pick_entries_of_the_full_rows(self):
         model = self._model()
         outputs = DFlashBaseModelOutput(None, base_hidden=self._hidden(model))
-        positions = self._positions()
-        token_ids = torch.randint(0, model.config.vocab_size, (*positions.shape, 5))
-        rows = model._teacher_logits(outputs, positions)
+        labels = self._labels()
+        token_ids = torch.randint(0, model.config.vocab_size, (*labels.shape, 5))
+        rows = model._teacher_logits(outputs, labels)
         torch.testing.assert_close(
-            model._teacher_logits(outputs, positions, token_ids), rows.gather(-1, token_ids)
+            model._teacher_logits(outputs, labels, token_ids), rows.gather(-1, token_ids)
         )
 
     def test_handed_over_logits_are_gathered_not_recomputed(self):
         model = self._model()
         logits = torch.randn(2, SEQ_LEN, model.config.vocab_size)
         outputs = DFlashBaseModelOutput(None, logits=logits)
-        positions = self._positions()
-        token_ids = torch.randint(0, model.config.vocab_size, (*positions.shape, 5))
-        want = self._rows(logits, positions)
-        assert torch.equal(model._teacher_logits(outputs, positions), want)
+        labels = self._labels()
+        token_ids = torch.randint(0, model.config.vocab_size, (*labels.shape, 5))
+        want = self._rows(logits, labels - 1)
+        assert torch.equal(model._teacher_logits(outputs, labels), want)
         assert torch.equal(
-            model._teacher_logits(outputs, positions, token_ids), want.gather(-1, token_ids)
+            model._teacher_logits(outputs, labels, token_ids), want.gather(-1, token_ids)
         )
 
+    def test_label_zero_does_not_wrap_to_the_last_row(self):
+        model = self._model()
+        logits = torch.randn(1, SEQ_LEN, model.config.vocab_size)
+        outputs = DFlashBaseModelOutput(None, logits=logits)
+        got = model._teacher_logits(outputs, torch.tensor([[0]]))
+        assert torch.equal(got[0, 0], logits[0, 0])
+
     def test_missing_base_distribution_raises(self):
         model = self._model()
         with pytest.raises(ValueError, match="base_model_hidden_states"):
-            model._teacher_logits(DFlashBaseModelOutput(None), self._positions())
+            model._teacher_logits(DFlashBaseModelOutput(None), self._labels())
 
     def test_offline_dict_keeps_the_hidden_unprojected(self):
         d = {

From 8363a6a5e077069731112a723a28f685da33bf38 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 07:41:35 +0000
Subject: [PATCH 02/15] refactor(speculative): one zero loss and one
 target-hidden pick for the DFlash family

Two pieces of the family's forwards were written out once per model.

- The step with nothing to supervise returns a zero loss that every trainable draft parameter
  feeds, so DDP's reduction still sees each of them. There were four copies in two forms:
  DFlash's indexed one element of each trainable parameter, while Domino's forward, DSpark's
  forward and DSpark's degenerate-loss branch summed every whole parameter, reading all draft
  weights to produce a zero. HFDFlashModel._draft_zero_loss now has the cheap form once.
- The concatenation of the base layers the draft's fc reads, hidden_states[lid + 1] for each
  target layer, appeared in the DFlash, DSpark and Domino forwards. It is now
  HFDFlashModel._target_hidden, which also carries the "+ 1" (hidden_states[0] is the
  embedding output) that modify()'s check on negative ids refers to.

Losses and gradients are unchanged; DSpark's DDP-coverage tests, which require a gradient on
every draft parameter for a batch with no valid anchor and for one with zero total weight,
pass on the shared helper.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../torch/speculative/plugins/hf_dflash.py    | 37 ++++++++++++-------
 .../torch/speculative/plugins/hf_domino.py    |  8 +---
 .../torch/speculative/plugins/hf_dspark.py    | 15 ++------
 3 files changed, 29 insertions(+), 31 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 7d642602768..c13cf1ab095 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -418,7 +418,7 @@ def modify(self, config):
                     "contains duplicates; each draft layer needs a distinct capture layer "
                     "(the streaming producer captures each base layer at most once)."
                 )
-            # forward() indexes hidden_states[lid + 1], where a negative id silently wraps.
+            # _target_hidden indexes hidden_states[lid + 1], where a negative id silently wraps.
             if min(user_target_layer_ids) < 0:
                 raise ValueError(
                     f"dflash_architecture_config.target_layer_ids {user_target_layer_ids} "
@@ -744,6 +744,24 @@ def _build_generate_swa_mask(self, ctx_len, bsz, dtype, device):
         attn_mask.masked_fill_(~keep, torch.finfo(dtype).min)
         return attn_mask
 
+    def _target_hidden(self, hidden_states):
+        """The base hidden states the draft's ``fc`` reads, concatenated: [B, seq, n * H].
+
+        ``hidden_states[0]`` is the embedding output, so layer ``lid``'s output is at ``lid + 1``.
+        """
+        return torch.cat([hidden_states[lid + 1] for lid in self.target_layer_ids], dim=-1)
+
+    def _draft_zero_loss(self, device):
+        """A zero loss that every trainable draft parameter feeds, for a step with nothing to train.
+
+        DDP expects a gradient for each parameter on every rank; a rank returning a constant
+        loss would leave the others waiting on its reduction.
+        """
+        return sum(
+            (p.reshape(-1)[0] * 0.0 for p in self.dflash_module.parameters() if p.requires_grad),
+            torch.zeros((), device=device),
+        )
+
     @torch.no_grad()
     def _teacher_logits(self, base_outputs, label_positions, token_ids=None):
         """The base model's logits for the tokens at ``label_positions`` ([B, ...]) -> [B, ..., vocab].
@@ -1007,9 +1025,7 @@ def forward(
                     "The base model did not return hidden states required for DFlash training. "
                     "Ensure its top-level multimodal forward supports output_hidden_states=True."
                 )
-            offset = 1
-            selected = [raw_outputs.hidden_states[lid + offset] for lid in self.target_layer_ids]
-            target_hidden = torch.cat(selected, dim=-1)  # [B, seq, num_layers * H]
+            target_hidden = self._target_hidden(raw_outputs.hidden_states)
             base_outputs = DFlashBaseModelOutput(
                 target_hidden=target_hidden, logits=raw_outputs.logits
             )
@@ -1035,17 +1051,10 @@ def forward(
         n_blocks = anchor_positions.shape[1]
 
         if n_blocks == 0 or not block_keep_mask.any():
-            # Keep all trainable draft parameters in the graph so DDP can reduce a rank
-            # that receives an all-masked answer-only batch.
-            dummy = sum(
-                (
-                    parameter.reshape(-1)[0] * 0.0
-                    for parameter in self.dflash_module.parameters()
-                    if parameter.requires_grad
-                ),
-                torch.zeros((), device=device),
+            # An all-masked answer-only batch: nothing to train, but DDP still reduces.
+            return ModelOutput(
+                loss=self._draft_zero_loss(device), logits=base_outputs.logits, train_acc=[[0.0]]
             )
-            return ModelOutput(loss=dummy, logits=base_outputs.logits, train_acc=[[0.0]])
 
         # 4. Build draft inputs
         noise_embedding = self._build_noise_embedding(
diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py
index 45d4e1be6c7..5cf502ecb88 100644
--- a/modelopt/torch/speculative/plugins/hf_domino.py
+++ b/modelopt/torch/speculative/plugins/hf_domino.py
@@ -323,9 +323,7 @@ def forward(
                     attention_mask=attention_mask,
                     output_hidden_states=True,
                 )
-            offset = 1
-            selected = [base_out.hidden_states[lid + offset] for lid in self.target_layer_ids]
-            target_hidden = torch.cat(selected, dim=-1)  # [B, seq, num_layers * H]
+            target_hidden = self._target_hidden(base_out.hidden_states)
 
         # 2. Build loss mask (same convention as DFlash).
         if labels is not None:
@@ -344,9 +342,7 @@ def forward(
         n_blocks = anchor_positions.shape[1]
 
         if n_blocks == 0 or not block_keep_mask.any():
-            # Zero loss that still flows through all draft params for DDP sync.
-            dummy = sum(p.sum() for p in self.dflash_module.parameters()) * 0.0
-            return ModelOutput(loss=dummy, logits=None, train_acc=[[0.0]])
+            return ModelOutput(loss=self._draft_zero_loss(device), logits=None, train_acc=[[0.0]])
 
         # 4. Build draft inputs (inherited helpers).
         noise_embedding = self._build_noise_embedding(
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index 58ae972d628..144b8d635fb 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -236,11 +236,8 @@ def _compute_dspark_loss(
         valid_count = flat_weights.sum() + 1e-6
 
         if valid_count <= 1.0:
-            # Touch every draft parameter, as forward()'s early return does, so DDP with
-            # find_unused_parameters=False still sees the confidence head's gradient.
-            loss = (
-                flat_final.sum() * 0.0 + sum(p.sum() for p in self.dflash_module.parameters()) * 0.0
-            )
+            # As forward()'s early return: DDP still reduces the confidence head's gradient.
+            loss = self._draft_zero_loss(flat_final.device)
             metrics = {"ce_loss": 0.0, "l1_loss": 0.0, "confidence_loss": 0.0, "base_accuracy": 0.0}
             return loss, 0.0, metrics
 
@@ -352,10 +349,8 @@ def forward(
                     attention_mask=attention_mask,
                     output_hidden_states=True,
                 )
-            offset = 1
-            selected = [base_out.hidden_states[lid + offset] for lid in self.target_layer_ids]
             base_outputs = DFlashBaseModelOutput(
-                target_hidden=torch.cat(selected, dim=-1),  # [B, seq, num_layers * H]
+                target_hidden=self._target_hidden(base_out.hidden_states),
                 base_hidden=base_out.last_hidden_state,
             )
         target_hidden = base_outputs.target_hidden
@@ -377,9 +372,7 @@ def forward(
         n_blocks = anchor_positions.shape[1]
 
         if n_blocks == 0 or not block_keep_mask.any():
-            # Zero loss that still flows through all draft params for DDP sync.
-            dummy = sum(p.sum() for p in self.dflash_module.parameters()) * 0.0
-            return ModelOutput(loss=dummy, logits=None, train_acc=[[0.0]])
+            return ModelOutput(loss=self._draft_zero_loss(device), logits=None, train_acc=[[0.0]])
 
         # 4. Build draft inputs (inherited helpers).
         noise_embedding = self._build_noise_embedding(

From 58a387470254a0b385aee6c48685ff615cb5319d Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 07:42:59 +0000
Subject: [PATCH 03/15] refactor(speculative): build the generation mask from
 the training rule

_build_generate_swa_mask hand-wrote the draft's visibility rule a second time -- context
before the block, the window counted from each query's real position, block-causal under
dflash_draft_attention="causal" -- and its docstring asked readers to keep it in step with
training. It is the training rule, _draft_mask_mod, for the one block being drafted,
anchored at ctx_len, so it now evaluates that. The grid evaluation and the additive
conversion that the training mask's dense path already did become _dense_additive_mask,
shared by both.

A verbatim copy of the old builder stays in the tests as the reference: over contexts of 1,
5 and 16 tokens, no window, a window of one block and of two, and bidirectional and causal
blocks, the masks are equal element for element (and None where the old one returned None).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../torch/speculative/plugins/hf_dflash.py    | 52 ++++++++-----------
 .../speculative/plugins/test_hf_dflash.py     | 43 +++++++++++++++
 2 files changed, 65 insertions(+), 30 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index c13cf1ab095..f3f190032a4 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -162,6 +162,19 @@ def mask_mod(b, h, q_idx, kv_idx):
     return mask_mod
 
 
+def _dense_additive_mask(mask_mod, bsz, q_len, kv_len, dtype, device):
+    """``mask_mod`` evaluated on broadcast index grids, as SDPA's additive [B, 1, Q, KV] mask."""
+    visible = mask_mod(
+        torch.arange(bsz, device=device).view(-1, 1, 1, 1),
+        None,
+        torch.arange(q_len, device=device).view(1, 1, -1, 1),
+        torch.arange(kv_len, device=device).view(1, 1, 1, -1),
+    )
+    attn_mask = torch.zeros(bsz, 1, q_len, kv_len, device=device, dtype=dtype)
+    attn_mask.masked_fill_(~visible, torch.finfo(dtype).min)
+    return attn_mask
+
+
 def _dpace_position_weights(
     confidences: torch.Tensor, alpha: float, valid_mask: torch.Tensor | None = None
 ) -> torch.Tensor:
@@ -702,47 +715,26 @@ def _build_draft_attention_mask(
                 head_dims={layer.self_attn.head_dim for layer in self.dflash_module.layers},
             )
 
-        visible = mask_mod(
-            torch.arange(bsz, device=device).view(-1, 1, 1, 1),
-            None,
-            torch.arange(q_len, device=device).view(1, 1, -1, 1),
-            torch.arange(kv_len, device=device).view(1, 1, 1, -1),
-        )  # [B, 1, Q, KV]
-        attn_mask = torch.zeros(bsz, 1, q_len, kv_len, device=device, dtype=dtype)
-        attn_mask.masked_fill_(~visible, torch.finfo(dtype).min)
-        return attn_mask
+        return _dense_additive_mask(mask_mod, bsz, q_len, kv_len, dtype, device)
 
     def _build_generate_swa_mask(self, ctx_len, bsz, dtype, device):
         """Generation-time mask [B, 1, block_size, ctx_len + block_size], or None.
 
         Returns None only when there is nothing to mask: full attention over the context
-        *and* bidirectional blocks (KV cache with no mask). With sliding-window attention,
-        each block query only sees context within ``dflash_swa_window_size`` tokens before
-        its real position (ctx_len + position-in-block), matching training and vLLM
-        inference; block kv is left un-windowed. With ``dflash_draft_attention="causal"``
-        the block is additionally lower-triangular, mirroring
-        :meth:`_build_draft_attention_mask` so generation matches training.
+        *and* bidirectional blocks (KV cache with no mask). Otherwise it is the training rule,
+        :func:`_draft_mask_mod`, for the one block being drafted, anchored at ``ctx_len``:
+        the window counts from each query's real position, and under
+        ``dflash_draft_attention="causal"`` the block is lower-triangular.
         """
         window = self.dflash_swa_window_size
         causal = self.dflash_draft_attention == "causal"
         if window is None and not causal:
             return None
         block_size = self.dflash_block_size
-        kv_len = ctx_len + block_size
-        kv_idx = torch.arange(kv_len, device=device).view(1, 1, 1, -1)
-        q_pos_in_block = torch.arange(block_size, device=device).view(1, 1, -1, 1)
-        q_real_pos = ctx_len + q_pos_in_block
-        is_ctx = kv_idx < ctx_len
-        # Context kv kept iff within the window (when windowing); block kv always visible
-        # unless the block is causal, in which case only positions <= the query's.
-        keep_ctx = is_ctx if window is None else (is_ctx & (kv_idx > q_real_pos - window))
-        keep_block = ~is_ctx
-        if causal:
-            keep_block = keep_block & ((kv_idx - ctx_len) <= q_pos_in_block)
-        keep = keep_ctx | keep_block
-        attn_mask = torch.zeros(bsz, 1, block_size, kv_len, device=device, dtype=dtype)
-        attn_mask.masked_fill_(~keep, torch.finfo(dtype).min)
-        return attn_mask
+        anchors = torch.full((bsz, 1), ctx_len, dtype=torch.long, device=device)
+        keep = torch.ones(bsz, 1, dtype=torch.bool, device=device)
+        mask_mod = _draft_mask_mod(ctx_len, anchors, keep, block_size, window, causal)
+        return _dense_additive_mask(mask_mod, bsz, block_size, ctx_len + block_size, dtype, device)
 
     def _target_hidden(self, hidden_states):
         """The base hidden states the draft's ``fc`` reads, concatenated: [B, seq, n * H].
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index ae8337daba5..c178dd025f3 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -562,6 +562,49 @@ def test_padded_query_rows_index_in_bounds(self):
         assert visible.shape == (padded_q.numel(), kv.numel())
 
 
+def _legacy_generate_swa_mask(model, ctx_len, bsz, dtype, device):
+    """Verbatim copy of the hand-written generation mask, kept as the semantic reference."""
+    window = model.dflash_swa_window_size
+    causal = model.dflash_draft_attention == "causal"
+    if window is None and not causal:
+        return None
+    block_size = model.dflash_block_size
+    kv_len = ctx_len + block_size
+    kv_idx = torch.arange(kv_len, device=device).view(1, 1, 1, -1)
+    q_pos_in_block = torch.arange(block_size, device=device).view(1, 1, -1, 1)
+    q_real_pos = ctx_len + q_pos_in_block
+    is_ctx = kv_idx < ctx_len
+    keep_ctx = is_ctx if window is None else (is_ctx & (kv_idx > q_real_pos - window))
+    keep_block = ~is_ctx
+    if causal:
+        keep_block = keep_block & ((kv_idx - ctx_len) <= q_pos_in_block)
+    keep = keep_ctx | keep_block
+    attn_mask = torch.zeros(bsz, 1, block_size, kv_len, device=device, dtype=dtype)
+    attn_mask.masked_fill_(~keep, torch.finfo(dtype).min)
+    return attn_mask
+
+
+class TestGenerateMaskRule:
+    """Generation evaluates the training rule for its one block, anchored at ``ctx_len``."""
+
+    @pytest.mark.parametrize("ctx_len", [1, 5, 16])
+    @pytest.mark.parametrize("window", [None, BLOCK_SIZE, 2 * BLOCK_SIZE])
+    @pytest.mark.parametrize("attention", ["bidirectional", "causal"])
+    def test_matches_the_hand_written_mask(self, ctx_len, window, attention):
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = get_dflash_config(block_size=BLOCK_SIZE)
+        config["dflash_swa_window_size"] = window
+        config["dflash_draft_attention"] = attention
+        mtsp.convert(model, [("dflash", config)])
+        args = (ctx_len, 2, torch.float32, torch.device("cpu"))
+        got = model._build_generate_swa_mask(*args)
+        want = _legacy_generate_swa_mask(model, *args)
+        if want is None:
+            assert got is None
+        else:
+            assert torch.equal(got, want)
+
+
 class TestDFlashSwaMask:
     """Test all-layer non-causal sliding-window attention mask (MiMo-style)."""
 

From 3c7ce77b1a328848d3d57c69e05c7e8cf89b976f Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 11:41:54 +0000
Subject: [PATCH 04/15] refactor(speculative): one block-supervision rule for
 the DFlash family

Which token each block position is trained on, which positions count, and how they are
weighted were written out five times: DFlash _compute_loss, DFlash2 _compute_loss (for the
selector), LiLiCorr _block_targets, DSpark _compute_dspark_loss and Domino
_compute_domino_loss. The copies differed only in the label offset (DFlash, DFlash2 and
LiLiCorr train position k on the token at anchor + k; DSpark and Domino on anchor + k + 1)
and in whether slot 0 is skipped, so a fix to one never reached the others.

- _block_targets returns the label token, its sequence position and the supervision mask,
  given the class constant _label_shift (0, or 1 for DSpark and Domino). The teacher reads
  its rows from those positions, including LiLiCorr's candidate logits.
- _block_tokens is the teacher-forced token gather that DSpark's and Domino's heads and
  DFlash2's predecessor ids each wrote by hand.
- _block_position_weights holds the D-PACE and decay weighting. DSpark and Domino pass
  "decay" explicitly, which is what they do today: they never read dflash_loss_objective.
- _masked_top1 replaces five copies of the masked top-1 accuracy.

The in-loss zero branches keep routing through the draft's output (flat_logits.sum() * 0.0)
rather than _draft_zero_loss: with DFlashModule as its own FSDP2 unit (the MiniMax-M2.7
launcher layout), reading the resharded parameters directly skips the unit's backward
collectives on that rank and the job aborts with a collective mismatch.

Behavior is unchanged. Against the previous head, a train step of DFlash (KD and CE),
DFlash2 (ce, tv and lambda, with the selector), LiLiCorr, DSpark (vanilla and rnn heads,
confidence head on) and Domino, offline and online, with D-PACE and decay, and steps that
take the zero-loss branch, gives bitwise-equal loss, gradients and metrics in fp32 and bf16.
DFlash2's predecessor ids now come from the shared gather; they differ from the old
formula only at positions whose label lies past the sequence end, which carry no selector
weight.

A new test pins each variant's labels and supervision on a hand-worked example; it fails if
DSpark's shift is set to 0 or if the anchor slot stops being excluded. No existing test did.

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_dflash.py    | 147 +++++++++++-------
 .../torch/speculative/plugins/hf_dflash2.py   |  38 ++---
 .../torch/speculative/plugins/hf_domino.py    |  71 ++++-----
 .../torch/speculative/plugins/hf_dspark.py    |  72 ++++-----
 .../torch/speculative/plugins/hf_lilicorr.py  |  67 +++-----
 .../speculative/plugins/test_hf_dflash.py     |  41 +++++
 6 files changed, 227 insertions(+), 209 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index f3f190032a4..08ad5699120 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -133,6 +133,78 @@ def _gather_rows(x, positions):
     return torch.gather(x, 1, idx).reshape(*positions.shape, x.size(-1))
 
 
+def _block_positions(anchor_positions, block_size, offset, seq_len):
+    """Index ``anchor + offset + k`` of each block position ``k``, clamped onto the sequence.
+
+    Returns ``(positions [B, N, bs], inside)``; ``inside`` is False past the sequence end.
+    """
+    k = torch.arange(block_size, device=anchor_positions.device).view(1, 1, -1)
+    unclamped = anchor_positions.unsqueeze(-1) + k + offset
+    return unclamped.clamp(min=0, max=seq_len - 1), unclamped < seq_len
+
+
+def _block_gather(x, positions):
+    """``x[b, positions[b, n, k]]`` for ``x`` [B, seq] and ``positions`` [B, N, bs]."""
+    return torch.gather(x.unsqueeze(1).expand(-1, positions.shape[1], -1), 2, positions)
+
+
+def _block_tokens(input_ids, anchor_positions, block_size, offset):
+    """The token at ``anchor + offset + k`` for each block position ``k``: [B, N, bs]."""
+    positions, _ = _block_positions(anchor_positions, block_size, offset, input_ids.shape[1])
+    return _block_gather(input_ids, positions)
+
+
+class _BlockTargets(NamedTuple):
+    """What each block position is trained on: the token at ``anchor + k + label_shift``."""
+
+    token_ids: torch.Tensor  # [B, N, bs]
+    positions: torch.Tensor  # [B, N, bs] their sequence indices, clamped onto the sequence
+    mask: torch.Tensor  # [B, N, bs] 1.0 where the position is supervised
+
+
+def _block_targets(input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, shift):
+    """Labels of every block position and where they are supervised.
+
+    A position is supervised when its block is kept and its label lies inside the sequence
+    and the loss mask, and is not the anchor itself, which the draft is given as input.
+    """
+    positions, inside = _block_positions(anchor_positions, block_size, shift, input_ids.shape[1])
+    k = torch.arange(block_size, device=input_ids.device).view(1, 1, -1)
+    mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float()
+    mask = mask * inside.float()
+    mask = mask * (k + shift > 0).float()
+    mask = mask * _block_gather(loss_mask, positions)
+    return _BlockTargets(_block_gather(input_ids, positions), positions, mask)
+
+
+def _block_position_weights(mask, scheme, first, decay_factor, dpace_alpha, target_ce=None):
+    """``mask`` weighted per block position, by D-PACE or by exponential decay.
+
+    ``first`` is the first predicted block position. ``target_ce`` returns the draft's
+    cross-entropy on its labels ([B, N, bs]); only D-PACE calls it.
+    """
+    block_size = mask.shape[-1]
+    if scheme == "dpace" and block_size > 1:
+        with torch.no_grad():
+            # Draft confidence q_i = exp(-CE) on the label; detached, as in the paper (Eq. 9).
+            confidences = torch.exp(-target_ce()[..., first:].float())
+            dpace = torch.ones_like(mask)
+            dpace[..., first:] = _dpace_position_weights(
+                confidences, dpace_alpha, valid_mask=mask[..., first:]
+            )
+        return mask * dpace
+    if decay_factor > 0:
+        k = torch.arange(block_size, device=mask.device).view(1, 1, -1)
+        return mask * torch.exp(-(k - first).clamp(min=0).float() / decay_factor)
+    return mask
+
+
+def _masked_top1(logits, target_ids, mask):
+    """Top-1 accuracy of ``logits`` against ``target_ids`` where ``mask`` is set (0-dim)."""
+    correct = (logits.argmax(dim=-1) == target_ids) & (mask > 0.5)
+    return correct.sum().float() / (mask.sum() + 1e-6)
+
+
 def _draft_mask_mod(seq_len, anchor_positions, block_keep_mask, block_size, window, causal):
     """The draft's attention visibility rule, as a FlexAttention ``mask_mod``.
 
@@ -243,6 +315,10 @@ class _DFlashLossTerms(NamedTuple):
 class HFDFlashModel(DFlashModel):
     """DFlash Model for HuggingFace transformers."""
 
+    # Block position k is trained on the token at anchor + k + _label_shift: DFlash predicts
+    # the anchor's own position onward, so slot 0 (the given anchor) is never supervised.
+    _label_shift = 0
+
     @property
     def _base_model(self):
         return self.get_submodule(self.base_model_path)
@@ -819,35 +895,14 @@ def _compute_loss(
         Returns:
             (loss, accuracy) tuple.
         """
-        bsz, seq_len = input_ids.shape
         block_size = self.dflash_block_size
-        n_blocks = anchor_positions.shape[1]
-        device = input_ids.device
-
-        label_offsets = torch.arange(0, block_size, device=device).view(1, 1, -1)
-        label_indices = anchor_positions.unsqueeze(-1) + label_offsets
-        valid_label = label_indices < seq_len
-        safe_label_indices = label_indices.clamp(max=seq_len - 1)
-
-        target_ids = torch.gather(
-            input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
-        )
-
-        # Weight mask: valid block * in bounds * exclude anchor (pos 0) * loss_mask
-        weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float()
-        weight_mask = weight_mask * valid_label.float()
-        pos_in_block = torch.arange(block_size, device=device).view(1, 1, -1)
-        weight_mask = weight_mask * (pos_in_block > 0).float()
-
-        orig_loss_mask = torch.gather(
-            loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
+        targets = _block_targets(
+            input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, self._label_shift
         )
-        weight_mask = weight_mask * orig_loss_mask
-
-        binary_eval_mask = weight_mask.view(-1)
+        binary_eval_mask = targets.mask.view(-1)
 
         flat_logits = logits.view(-1, logits.size(-1))
-        flat_targets = target_ids.view(-1)
+        flat_targets = targets.token_ids.view(-1)
 
         kd = self.dflash_self_logit_distillation
         if kd and base_outputs is None:
@@ -862,35 +917,25 @@ def _compute_loss(
         if not kd:
             loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none")
 
-        # Block-position loss weighting: dynamic D-PACE weights or static exponential decay.
-        if self.dflash_loss_objective == "dpace" and block_size > 1:
-            # Draft confidence q_i = exp(-CE) on the target-selected token, over the
-            # predicted positions (slot 0 is the given anchor, already masked above).
-            # Weights are detached (paper Eq.9), so this adds the documented ~2.3%
-            # training overhead without altering the cross-entropy gradient.
-            with torch.no_grad():
-                conf_ce = (
-                    loss_per_token.detach()
-                    if loss_per_token is not None
-                    else F.cross_entropy(flat_logits, flat_targets, reduction="none")
-                ).view(bsz, n_blocks, block_size)
-                confidences = torch.exp(-conf_ce[..., 1:].float())
-                dpace = torch.ones_like(weight_mask)
-                dpace[..., 1:] = _dpace_position_weights(
-                    confidences, self.dflash_dpace_alpha, valid_mask=weight_mask[..., 1:]
-                )
-            weight_mask = weight_mask * dpace
-        elif self.dflash_loss_decay_factor > 0:
-            k = torch.arange(block_size, device=device).view(1, 1, -1)
-            decay = torch.exp(-(k - 1).clamp(min=0).float() / self.dflash_loss_decay_factor)
-            weight_mask = weight_mask * decay
+        weight_mask = _block_position_weights(
+            targets.mask,
+            self.dflash_loss_objective,
+            1 - self._label_shift,
+            self.dflash_loss_decay_factor,
+            self.dflash_dpace_alpha,
+            target_ce=lambda: (
+                loss_per_token.detach()
+                if loss_per_token is not None
+                else F.cross_entropy(flat_logits, flat_targets, reduction="none")
+            ).view(targets.mask.shape),
+        )
 
         flat_weights = weight_mask.view(-1)
         valid_count = flat_weights.sum() + 1e-6
 
         if valid_count > 1.0:
             if kd:
-                flat_teacher = self._teacher_logits(base_outputs, safe_label_indices)
+                flat_teacher = self._teacher_logits(base_outputs, targets.positions)
                 flat_teacher = flat_teacher.reshape(flat_logits.shape)
                 target_soft = torch.softmax(flat_teacher, dim=-1)
                 draft_logsoft = torch.log_softmax(flat_logits, dim=-1)
@@ -900,11 +945,9 @@ def _compute_loss(
                 loss = (loss_per_token * flat_weights).sum() / valid_count
 
             with torch.no_grad():
-                preds = flat_logits.argmax(dim=-1)
-                correct = (preds == flat_targets) & (binary_eval_mask > 0.5)
-                accuracy = correct.sum().float() / (binary_eval_mask.sum() + 1e-6)
-                accuracy = accuracy.item()
+                accuracy = _masked_top1(flat_logits, flat_targets, binary_eval_mask).item()
         else:
+            # Through the draft's output, so an FSDP unit around the draft still runs its backward.
             loss = flat_logits.sum() * 0.0
             accuracy = 0.0
 
diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py
index 6060923b546..5ddc95a6b48 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash2.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash2.py
@@ -39,7 +39,7 @@
 from transformers import PreTrainedModel
 
 from ..dflash.conversion import DFlash2DMRegistry
-from .hf_dflash import HFDFlashModel
+from .hf_dflash import HFDFlashModel, _block_targets, _block_tokens
 from .modeling_dflash2 import DFlash2Module
 
 __all__ = ["HFDFlash2Model"]
@@ -231,45 +231,35 @@ def _compute_loss(
         if self.dflash_selector_loss_alpha <= 0 or draft_hidden is None:
             return loss, accuracy
 
-        bsz, seq_len = input_ids.shape
+        bsz = input_ids.shape[0]
         block_size = self.dflash_block_size
         n_blocks = anchor_positions.shape[1]
-        device = input_ids.device
-
-        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
+
+        # The backbone's supervision mask, without its position weighting (decay/D-PACE):
+        # that shapes *where* the backbone spends capacity, while the selector should learn
+        # every position's transition equally.
+        targets = _block_targets(
+            input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, self._label_shift
         )
 
         # 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 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.
         #
         # This is the train/serve contract: CandidateSelector.greedy_path seeds its walk
         # from the anchor token, so ITS position 0 corresponds to block offset 1. A caller
         # 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))
+        predecessor_ids = _block_tokens(
+            input_ids, anchor_positions, block_size, self._label_shift - 1
+        )
 
         selector_loss, selector_accuracy, selector_coverage = self._selector_loss(
             logits.reshape(bsz, n_blocks, block_size, -1),
-            target_ids,
+            targets.token_ids,
             draft_hidden.reshape(bsz, n_blocks, block_size, -1),
             predecessor_ids,
-            weight_mask,
+            targets.mask,
         )
         self._selector_metrics = {
             "selector_accuracy": selector_accuracy,
diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py
index 5cf502ecb88..91507974830 100644
--- a/modelopt/torch/speculative/plugins/hf_domino.py
+++ b/modelopt/torch/speculative/plugins/hf_domino.py
@@ -68,7 +68,13 @@
 from transformers.utils import ModelOutput
 
 from ..dflash.conversion import DominoDMRegistry
-from .hf_dflash import HFDFlashModel
+from .hf_dflash import (
+    HFDFlashModel,
+    _block_position_weights,
+    _block_targets,
+    _block_tokens,
+    _masked_top1,
+)
 from .modeling_dflash import DFlashBaseModelOutput
 from .modeling_domino import DominoModule
 
@@ -101,6 +107,8 @@ class HFDominoModel(HFDFlashModel):
     route to it when ``dflash_architecture_config.projector_type == "domino"``.
     """
 
+    _label_shift = 1
+
     def _build_draft_module(self, dflash_config):
         """Build the Domino draft module (DFlash backbone + GRU correction head)."""
         return DominoModule(dflash_config)
@@ -151,18 +159,15 @@ def _apply_domino_head(self, hidden, base_logits, input_ids, anchor_positions, n
         Returns:
             Corrected logits [B, N*block_size, vocab].
         """
-        bsz, seq_len = input_ids.shape
+        bsz = input_ids.shape[0]
         bs = self.dflash_block_size
-        device = input_ids.device
         suffix_start = self.dflash_module.pure_draft_prefix_len
 
         hidden4d = hidden.reshape(bsz, n_blocks, bs, hidden.size(-1))
         base4d = base_logits.reshape(bsz, n_blocks, bs, -1)
 
         # Teacher-forced previous tokens: the real token at anchor+j for j in [0, bs).
-        prev_offsets = torch.arange(bs, device=device).view(1, 1, -1)
-        prev_idx = (anchor_positions.unsqueeze(-1) + prev_offsets).clamp(max=seq_len - 1)
-        prev_ids = torch.gather(input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, prev_idx)
+        prev_ids = _block_tokens(input_ids, anchor_positions, bs, self._label_shift - 1)
 
         block_emb = self._base_model_embeddings(prev_ids)  # [B, N, bs, H]
         gru_in = block_emb.reshape(bsz * n_blocks, bs, block_emb.size(-1))
@@ -187,40 +192,27 @@ def _compute_domino_loss(
         Uses next-token (shift_label) alignment: position k predicts the token at
         anchor+k+1, and position 0 is *not* excluded (unlike base DFlash).
         """
-        bsz, seq_len = input_ids.shape
-        bs = self.dflash_block_size
-        n_blocks = anchor_positions.shape[1]
-        device = input_ids.device
-
-        # shift_label=True: label for block position k is the token at anchor+k+1.
-        label_offsets = torch.arange(1, 1 + bs, device=device).view(1, 1, -1)
-        label_indices = anchor_positions.unsqueeze(-1) + label_offsets
-        valid_label = label_indices < seq_len
-        safe_label_indices = label_indices.clamp(max=seq_len - 1)
-
-        target_ids = torch.gather(
-            input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
+        targets = _block_targets(
+            input_ids,
+            anchor_positions,
+            block_keep_mask,
+            loss_mask,
+            self.dflash_block_size,
+            self._label_shift,
         )
-
-        # Weight mask: valid block * in bounds * loss_mask. No pos-0 exclusion.
-        weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, bs).float()
-        weight_mask = weight_mask * valid_label.float()
-        orig_loss_mask = torch.gather(
-            loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
+        binary_eval_mask = targets.mask.view(-1)
+        # dflash_loss_objective is not wired for Domino yet: it always decays.
+        weight_mask = _block_position_weights(
+            targets.mask,
+            "decay",
+            1 - self._label_shift,
+            self.dflash_loss_decay_factor,
+            self.dflash_dpace_alpha,
         )
-        weight_mask = weight_mask * orig_loss_mask
-
-        binary_eval_mask = weight_mask.view(-1)
-
-        # Loss decay: exp(-k/gamma) so the first prediction (k=0) gets weight 1.0.
-        if self.dflash_loss_decay_factor > 0:
-            k = torch.arange(bs, device=device).view(1, 1, -1)
-            decay = torch.exp(-k.clamp(min=0).float() / self.dflash_loss_decay_factor)
-            weight_mask = weight_mask * decay
 
         flat_final = final_logits.reshape(-1, final_logits.size(-1))
         flat_base = base_logits.reshape(-1, base_logits.size(-1))
-        flat_targets = target_ids.reshape(-1)
+        flat_targets = targets.token_ids.reshape(-1)
         flat_weights = weight_mask.reshape(-1)
         valid_count = flat_weights.sum() + 1e-6
 
@@ -236,13 +228,8 @@ def _compute_domino_loss(
             loss = (1.0 - lambda_base) * final_loss + lambda_base * base_loss
 
             with torch.no_grad():
-                eval_count = binary_eval_mask.sum() + 1e-6
-                final_correct = (flat_final.argmax(dim=-1) == flat_targets) & (
-                    binary_eval_mask > 0.5
-                )
-                base_correct = (flat_base.argmax(dim=-1) == flat_targets) & (binary_eval_mask > 0.5)
-                accuracy = (final_correct.sum().float() / eval_count).item()
-                base_accuracy = (base_correct.sum().float() / eval_count).item()
+                accuracy = _masked_top1(flat_final, flat_targets, binary_eval_mask).item()
+                base_accuracy = _masked_top1(flat_base, flat_targets, binary_eval_mask).item()
         else:
             loss = flat_final.sum() * 0.0
             final_loss = loss
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index 144b8d635fb..924fb9a9cd7 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -69,7 +69,13 @@
 from transformers.utils import ModelOutput
 
 from ..dflash.conversion import DSparkDMRegistry
-from .hf_dflash import HFDFlashModel
+from .hf_dflash import (
+    HFDFlashModel,
+    _block_position_weights,
+    _block_targets,
+    _block_tokens,
+    _masked_top1,
+)
 from .modeling_dflash import DFlashBaseModelOutput
 from .modeling_dspark import DSparkModule
 
@@ -117,6 +123,8 @@ class HFDSparkModel(HFDFlashModel):
     to it when ``dflash_architecture_config.projector_type == "dspark"``.
     """
 
+    _label_shift = 1
+
     def _build_draft_module(self, dflash_config):
         """Build the DSpark draft module (DFlash backbone + Markov/confidence head)."""
         return DSparkModule(dflash_config)
@@ -158,18 +166,15 @@ def _apply_markov_head(self, hidden, backbone_logits, input_ids, anchor_position
 
         Returns ``(final_logits [B, N, bs, V], confidence_logits [B, N, bs] | None)``.
         """
-        bsz, seq_len = input_ids.shape
+        bsz = input_ids.shape[0]
         bs = self.dflash_block_size
-        device = input_ids.device
 
         hidden4d = hidden.reshape(bsz, n_blocks, bs, hidden.size(-1))
         base4d = backbone_logits.reshape(bsz, n_blocks, bs, -1)
 
         # Teacher-forced previous token for block position k: the real token at
         # anchor+k (so position 0's predecessor is the anchor itself).
-        prev_offsets = torch.arange(bs, device=device).view(1, 1, -1)
-        prev_idx = (anchor_positions.unsqueeze(-1) + prev_offsets).clamp(max=seq_len - 1)
-        prev_ids = torch.gather(input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, prev_idx)
+        prev_ids = _block_tokens(input_ids, anchor_positions, bs, self._label_shift - 1)
 
         bias = self.dflash_module.compute_markov_bias(prev_ids, hidden4d)
         final4d = base4d + bias
@@ -197,41 +202,28 @@ def _compute_dspark_loss(
         next-token distribution at position anchor+k (= label index - 1), read from
         ``base_outputs``.
         """
-        bsz, seq_len = input_ids.shape
-        bs = self.dflash_block_size
-        n_blocks = anchor_positions.shape[1]
-        device = input_ids.device
         vocab = final_logits.size(-1)
-
-        # shift_label=True: label for block position k is the token at anchor+k+1.
-        label_offsets = torch.arange(1, 1 + bs, device=device).view(1, 1, -1)
-        label_indices = anchor_positions.unsqueeze(-1) + label_offsets
-        valid_label = label_indices < seq_len
-        safe_label_indices = label_indices.clamp(max=seq_len - 1)
-
-        target_ids = torch.gather(
-            input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
+        targets = _block_targets(
+            input_ids,
+            anchor_positions,
+            block_keep_mask,
+            loss_mask,
+            self.dflash_block_size,
+            self._label_shift,
         )
-
-        # Weight mask: valid block * in bounds * loss_mask (no pos-0 exclusion).
-        weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, bs).float()
-        weight_mask = weight_mask * valid_label.float()
-        orig_loss_mask = torch.gather(
-            loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
+        binary_eval_mask = targets.mask.view(-1)
+        # dflash_loss_objective is not wired for DSpark yet: it always decays.
+        weight_mask = _block_position_weights(
+            targets.mask,
+            "decay",
+            1 - self._label_shift,
+            self.dflash_loss_decay_factor,
+            self.dflash_dpace_alpha,
         )
-        weight_mask = weight_mask * orig_loss_mask
-
-        binary_eval_mask = weight_mask.view(-1)
-
-        # Exponential position decay (exp(-k/gamma); position 0 gets weight 1).
-        if self.dflash_loss_decay_factor > 0:
-            k = torch.arange(bs, device=device).view(1, 1, -1)
-            decay = torch.exp(-k.float() / self.dflash_loss_decay_factor)
-            weight_mask = weight_mask * decay
 
         flat_final = final_logits.reshape(-1, vocab)
         flat_base = backbone_logits.reshape(-1, vocab)
-        flat_targets = target_ids.reshape(-1)
+        flat_targets = targets.token_ids.reshape(-1)
         flat_weights = weight_mask.reshape(-1)
         valid_count = flat_weights.sum() + 1e-6
 
@@ -242,7 +234,7 @@ def _compute_dspark_loss(
             return loss, 0.0, metrics
 
         # Aligned target distribution: the base model's own prediction of each label.
-        flat_teacher = self._teacher_logits(base_outputs, safe_label_indices).reshape(-1, vocab)
+        flat_teacher = self._teacher_logits(base_outputs, targets.positions).reshape(-1, vocab)
 
         # Term 1: cross-entropy on the corrected (final) logits.
         ce_per_token = F.cross_entropy(flat_final, flat_targets, reduction="none")
@@ -273,12 +265,8 @@ def _compute_dspark_loss(
         )
 
         with torch.no_grad():
-            eval_count = binary_eval_mask.sum() + 1e-6
-            keep = binary_eval_mask > 0.5
-            acc = ((flat_final.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count
-            base_acc = (
-                (flat_base.argmax(dim=-1) == flat_targets) & keep
-            ).sum().float() / eval_count
+            acc = _masked_top1(flat_final, flat_targets, binary_eval_mask)
+            base_acc = _masked_top1(flat_base, flat_targets, binary_eval_mask)
             # One device sync for all five scalars instead of one per .item().
             acc_v, base_acc_v, ce_v, l1_v, conf_v = torch.stack(
                 [acc, base_acc, ce_loss.detach(), l1_loss.detach(), confidence_loss.detach()]
diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py
index c77bbcd9f7d..be20a9d0f53 100644
--- a/modelopt/torch/speculative/plugins/hf_lilicorr.py
+++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py
@@ -77,7 +77,7 @@
 from transformers import PreTrainedModel
 
 from ..dflash.conversion import LiLiCorrDMRegistry
-from .hf_dflash import HFDFlashModel
+from .hf_dflash import HFDFlashModel, _block_positions, _block_targets
 from .modeling_lilicorr import LiLiCorrModule
 
 logger = logging.getLogger(__name__)
@@ -190,46 +190,6 @@ def get_exporter(self):
 
         return LiLiCorrExporter(self)
 
-    def _block_targets(self, input_ids, anchor_positions, block_keep_mask, loss_mask):
-        """Return ``(target_ids [B, N, bs], slot_mask [B, N, bs-1])`` for the lattice.
-
-        ``target_ids[b, n, j]`` is the ground-truth **vocabulary id** at absolute position
-        ``anchor[b, n] + j``, covering the whole block including position 0 (the anchor
-        token, which the caller slices off). Distinct from ``gt_indices`` in
-        :meth:`_compute_lilicorr_loss`, which is a position *within* a slot's k candidates.
-
-        ``slot_mask[b, n, s]`` says whether predicted slot ``s`` — block position
-        ``s + 1`` — is supervised: its block is kept, its position is in bounds, and
-        ``loss_mask`` allows it.
-
-        Recomputed rather than threaded out of :meth:`HFDFlashModel._compute_loss`: both
-        are exact functions of the anchors and the masks, so recomputing them keeps the
-        shared method's signature to the tensors a variant cannot derive on its own.
-        The slot mask is the parent's *pre-weighting* eval mask, so the two losses
-        supervise the same positions regardless of the block-position weighting.
-        """
-        bsz, seq_len = input_ids.shape
-        block_size = self.dflash_block_size
-        n_blocks = anchor_positions.shape[1]
-        device = input_ids.device
-
-        label_offsets = torch.arange(0, block_size, device=device).view(1, 1, -1)
-        label_indices = anchor_positions.unsqueeze(-1) + label_offsets
-        valid_label = label_indices < seq_len
-        safe_label_indices = label_indices.clamp(max=seq_len - 1)
-
-        target_ids = torch.gather(
-            input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
-        )
-        keep = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float()
-        keep = keep * valid_label.float()
-        keep = keep * torch.gather(
-            loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
-        )
-        # Block position 0 is the anchor — the one row whose input embedding is a real
-        # verified token — so the lattice covers positions 1..bs-1.
-        return target_ids, keep[:, :, 1:] > 0.5
-
     def _compute_loss(
         self,
         logits,
@@ -268,13 +228,22 @@ def _compute_loss(
             loss_mask,
             base_outputs=base_outputs,
         )
-        target_ids, slot_mask = self._block_targets(
-            input_ids, anchor_positions, block_keep_mask, loss_mask
+        # The backbone's own supervision, before its position weighting, so both losses
+        # train the same positions.
+        targets = _block_targets(
+            input_ids,
+            anchor_positions,
+            block_keep_mask,
+            loss_mask,
+            self.dflash_block_size,
+            self._label_shift,
         )
         lilicorr_loss, accuracy, metrics = self._compute_lilicorr_loss(
             logits=logits,
-            target_ids=target_ids,
-            slot_mask=slot_mask,
+            target_ids=targets.token_ids,
+            # Block position 0 is the anchor -- the one row whose input embedding is a real
+            # verified token -- so the lattice covers positions 1..bs-1.
+            slot_mask=targets.mask[:, :, 1:] > 0.5,
             anchor_positions=anchor_positions,
             draft_hidden=draft_hidden,
             target_hidden=target_hidden,
@@ -545,10 +514,10 @@ def _candidate_target_logits(
         batch_blocks, num_slots, topk = candidate_ids.shape
         bsz, n_blocks = anchor_positions.shape
         seq_len = base_outputs.target_hidden.shape[1]
-        # Slot s scores the token at anchor+1+s. Only the [blocks, slots, k] candidate logits
-        # are computed, never whole vocabulary rows.
-        slot_offsets = torch.arange(1, num_slots + 1, device=candidate_ids.device)
-        labels = (anchor_positions.unsqueeze(-1) + slot_offsets).clamp(max=seq_len - 1)
+        # Slot s is block position s + 1, labelled like the backbone's. Only the
+        # [blocks, slots, k] candidate logits are computed, never whole vocabulary rows.
+        labels, _ = _block_positions(anchor_positions, num_slots + 1, self._label_shift, seq_len)
+        labels = labels[..., 1:]
         candidate_target_logits = self._teacher_logits(
             base_outputs,
             labels,
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index c178dd025f3..d100b4a8f79 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -44,9 +44,14 @@
     DFlashBaseModelOutput,
     DFlashModule,
     HFDFlashModel,
+    _block_targets,
     _dpace_position_weights,
     build_target_layer_ids,
 )
+from modelopt.torch.speculative.plugins.hf_dflash2 import HFDFlash2Model
+from modelopt.torch.speculative.plugins.hf_domino import HFDominoModel
+from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel
+from modelopt.torch.speculative.plugins.hf_lilicorr import HFLiLiCorrModel
 from modelopt.torch.speculative.plugins.master_weight_adamw import (
     MasterWeightAdamW,
     VerifyMasterWeightsCallback,
@@ -1488,6 +1493,42 @@ def test_trailing_padding_blocks_do_not_change_the_loss(self):
         assert padded_acc == pytest.approx(base_acc)
 
 
+# Anchors 1 and 4 of an 8-token sequence whose loss_mask is 0 only at position 3:
+# (label positions, supervision mask) per block position.
+_LABELS_AT_ANCHOR = ([[1, 2, 3, 4], [4, 5, 6, 7]], [[0, 1, 0, 1], [0, 1, 1, 1]])
+_LABELS_AFTER_ANCHOR = ([[2, 3, 4, 5], [5, 6, 7, 7]], [[1, 0, 1, 1], [1, 1, 1, 0]])
+
+
+@pytest.mark.parametrize(
+    ("model_cls", "expected"),
+    [
+        (HFDFlashModel, _LABELS_AT_ANCHOR),
+        (HFDFlash2Model, _LABELS_AT_ANCHOR),
+        (HFLiLiCorrModel, _LABELS_AT_ANCHOR),
+        (HFDSparkModel, _LABELS_AFTER_ANCHOR),
+        (HFDominoModel, _LABELS_AFTER_ANCHOR),
+    ],
+    ids=["dflash", "dflash2", "lilicorr", "dspark", "domino"],
+)
+def test_block_targets_follow_the_label_shift(model_cls, expected):
+    """Each block position's label and supervision; a wrong shift trains silently worse."""
+    positions, mask = expected
+    input_ids = 100 + torch.arange(8).unsqueeze(0)  # token id = 100 + its position
+    loss_mask = torch.ones(1, 8)
+    loss_mask[0, 3] = 0.0
+    targets = _block_targets(
+        input_ids,
+        torch.tensor([[1, 4]]),
+        torch.tensor([[True, True]]),
+        loss_mask,
+        4,
+        model_cls._label_shift,
+    )
+    assert torch.equal(targets.positions, torch.tensor([positions]))
+    assert torch.equal(targets.token_ids, 100 + torch.tensor([positions]))
+    assert torch.equal(targets.mask, torch.tensor([mask], dtype=torch.float32))
+
+
 class TestTeacherLogits:
     """The base distribution is projected only at the rows, and vocab entries, asked for.
 

From b053414284ecbe0566308a6713cf31fbc3001d3b Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 11:43:21 +0000
Subject: [PATCH 05/15] refactor(speculative): share the generation step
 between DFlash and DSpark

DFlash's and DSpark's pseudo_speculative_generate repeated the same first half: run the
base model on the prefix, take its greedy next token, gather the target layers onto the
draft's device, build the block (anchor then mask tokens), its position ids and mask, and
run the draft. Only the decoding after that differs (DFlash takes the argmax of positions
1..bs-1; DSpark walks its Markov head over the block).

- _generation_base returns the base model's next token and hidden states, and carries the
  offline check both copies made.
- _draft_block drafts the block that token anchors. Its position ids now come from
  _build_position_ids at anchor = ctx_len -- the training layout -- the way the generation
  mask already comes from the training mask rule. The comment it replaces described a
  uniform-block layout training has not used since the anchors were introduced.
- _target_hidden takes an optional device, so the last two hand-written
  hidden_states[lid + 1] picks go through it.

The docstring's note that generation runs with attention_mask=None was stale since the
generation mask was added and now says what is passed.

Generated tokens are unchanged: DFlash (full attention, sliding window, causal blocks),
DFlash2, LiLiCorr, DSpark and Domino online give bitwise-equal draft tokens in fp32 and bf16
against the previous commit.

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_dflash.py    | 125 +++++++++---------
 .../torch/speculative/plugins/hf_dspark.py    |  46 +------
 2 files changed, 66 insertions(+), 105 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 08ad5699120..766ffaad7a7 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -812,12 +812,67 @@ def _build_generate_swa_mask(self, ctx_len, bsz, dtype, device):
         mask_mod = _draft_mask_mod(ctx_len, anchors, keep, block_size, window, causal)
         return _dense_additive_mask(mask_mod, bsz, block_size, ctx_len + block_size, dtype, device)
 
-    def _target_hidden(self, hidden_states):
+    def _target_hidden(self, hidden_states, device=None):
         """The base hidden states the draft's ``fc`` reads, concatenated: [B, seq, n * H].
 
         ``hidden_states[0]`` is the embedding output, so layer ``lid``'s output is at ``lid + 1``.
+        ``device`` gathers the layers first, for a base model sharded across GPUs.
         """
-        return torch.cat([hidden_states[lid + 1] for lid in self.target_layer_ids], dim=-1)
+        selected = [hidden_states[lid + 1] for lid in self.target_layer_ids]
+        if device is not None:
+            selected = [h.to(device) for h in selected]
+        return torch.cat(selected, dim=-1)
+
+    def _generation_base(self, input_ids):
+        """The base model's greedy next token after ``input_ids`` ([B, 1]), and its hidden states."""
+        if self.dflash_offline:
+            raise RuntimeError(
+                "An offline DFlash-family model cannot run AR validation / "
+                "pseudo_speculative_generate — base model layers were deleted during offline "
+                "conversion to save memory. Reload the full base model before running AR validation."
+            )
+        # Call the base model's inner model directly (avoids DynamicModule dispatch)
+        model_output = self._base_model(input_ids=input_ids, output_hidden_states=True)
+        base_logits = self._base_model_lm_head(model_output.last_hidden_state)
+        base_token = base_logits[:, -1:, :].argmax(dim=-1).to(input_ids.device)
+        return base_token, model_output.hidden_states
+
+    def _draft_block(self, base_token, hidden_states):
+        """Draft hidden states [B, block_size, H] for the block that ``base_token`` anchors.
+
+        Re-runs no base layer: the context is ``hidden_states`` from :meth:`_generation_base`.
+        """
+        # target_layer_ids spans early and late layers, so under device_map="auto" these come
+        # off different GPUs and the cat fails. Everything the draft consumes is gathered
+        # onto the draft's own device -- the last base layer's, per modify() -- which is
+        # also not input_ids.device once the base model is sharded. All no-ops single-device.
+        device = self._base_device()
+        target_hidden = self._target_hidden(hidden_states, device=device)
+
+        block_size = self.dflash_block_size
+        bsz = base_token.shape[0]
+
+        # Block: first token is base_token (anchor), rest are mask
+        block_ids = torch.full(
+            (bsz, block_size), self.mask_token_id, dtype=torch.long, device=device
+        )
+        block_ids[:, 0] = base_token.squeeze(-1).to(device)
+        # The embedding table is on the first shard; its output follows it, not block_ids.
+        noise_embedding = self._base_model_embeddings(block_ids).to(device)
+
+        # The training layout for one block anchored right after the context.
+        ctx_len = target_hidden.shape[1]
+        anchors = torch.full((bsz, 1), ctx_len, dtype=torch.long, device=device)
+        pos_ids = self._build_position_ids(ctx_len, anchors, device)
+
+        attn_mask = self._build_generate_swa_mask(ctx_len, bsz, target_hidden.dtype, device)
+
+        return self.dflash_module(
+            noise_embedding=noise_embedding,
+            target_hidden=target_hidden,
+            position_ids=pos_ids,
+            attention_mask=attn_mask,
+        )
 
     def _draft_zero_loss(self, device):
         """A zero loss that every trainable draft parameter feeds, for a step with nothing to train.
@@ -1160,7 +1215,7 @@ def pseudo_speculative_generate(self, input_ids, steps=1):
            - ``noise_embedding``: embedded block tokens
            - ``target_hidden``: extracted hidden states from step 2
            - ``position_ids``: context + block positions
-           - ``attention_mask=None``: no mask at inference (all positions attend freely)
+           - ``attention_mask``: the training rule for this one block (None when it masks nothing)
            The draft model's KV injection concatenates projected target_hidden as K/V
            with the block's own K/V, enabling the draft to "see" the target's context.
 
@@ -1184,68 +1239,10 @@ def pseudo_speculative_generate(self, input_ids, steps=1):
             base_token: Next token from base model [B, 1].
             draft_tokens: Draft tokens [B, min(steps, block_size-1)] or None if steps < 1.
         """
-        if self.dflash_offline:
-            raise RuntimeError(
-                "DFlash offline model cannot run AR validation / pseudo_speculative_generate — "
-                "base model layers were deleted during offline conversion to save memory. "
-                "Reload the full base model before running AR validation."
-            )
-        # Call the base model's inner model directly (avoids DynamicModule dispatch)
-        model_output = self._base_model(
-            input_ids=input_ids,
-            output_hidden_states=True,
-        )
-        # Compute logits via lm_head
-        base_logits = self._base_model_lm_head(model_output.last_hidden_state)
-        # Build output with hidden_states
-        base_outputs = ModelOutput(
-            logits=base_logits,
-            hidden_states=model_output.hidden_states,
-        )
-        base_logits = base_outputs.logits
-        base_token = base_logits[:, -1:, :].argmax(dim=-1).to(input_ids.device)
-
+        base_token, hidden_states = self._generation_base(input_ids)
         if steps < 1:
             return base_token, None
-
-        # Extract target hidden states (raw, before FC projection)
-        hid_offset = 1
-        selected = [base_outputs.hidden_states[lid + hid_offset] for lid in self.target_layer_ids]
-        # target_layer_ids spans early and late layers, so under device_map="auto" these come
-        # off different GPUs and the cat below fails. Everything the draft consumes is gathered
-        # onto the draft's own device -- the last base layer's, per modify() -- which is
-        # also not input_ids.device once the base model is sharded. All no-ops single-device.
-        device = self._base_device()
-        target_hidden = torch.cat([h.to(device) for h in selected], dim=-1)
-
-        block_size = self.dflash_block_size
-        bsz = input_ids.shape[0]
-
-        # Block: first token is base_token (anchor), rest are mask
-        block_ids = torch.full(
-            (bsz, block_size), self.mask_token_id, dtype=torch.long, device=device
-        )
-        block_ids[:, 0] = base_token.squeeze(-1).to(device)
-        # The embedding table is on the first shard; its output follows it, not block_ids.
-        noise_embedding = self._base_model_embeddings(block_ids).to(device)
-
-        # Position IDs: training uses [0..L-1, 0..L-1] where noise positions
-        # mirror context positions. At inference, block predicts tokens at
-        # seq_len..seq_len+B-1, so noise positions continue from ctx_len.
-        ctx_len = target_hidden.shape[1]
-        ctx_positions = torch.arange(ctx_len, device=device)
-        block_positions = torch.arange(ctx_len, ctx_len + block_size, device=device)
-        pos_ids = torch.cat([ctx_positions, block_positions]).unsqueeze(0).expand(bsz, -1)
-
-        attn_mask = self._build_generate_swa_mask(ctx_len, bsz, target_hidden.dtype, device)
-
-        # Draft forward
-        draft_hidden = self.dflash_module(
-            noise_embedding=noise_embedding,
-            target_hidden=target_hidden,
-            position_ids=pos_ids,
-            attention_mask=attn_mask,
-        )
+        draft_hidden = self._draft_block(base_token, hidden_states)
 
         # Logits on positions 1..block_size-1 (skip anchor at position 0)
         draft_logits = self._base_model_lm_head(draft_hidden[:, 1:, :])
@@ -1253,5 +1250,5 @@ def pseudo_speculative_generate(self, input_ids, steps=1):
 
         # Return up to `steps` tokens. base_token already follows input_ids; draft_tokens is
         # produced on the draft's device, and callers cat both onto the running sequence.
-        num_tokens = min(steps, block_size - 1)
+        num_tokens = min(steps, self.dflash_block_size - 1)
         return base_token, draft_tokens[:, :num_tokens].to(input_ids.device)
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index 924fb9a9cd7..d000a3efa10 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -415,53 +415,17 @@ def pseudo_speculative_generate(self, input_ids, steps=1):
         (the anchor for k=0). Without this override DSpark would fall back to the
         DFlash backbone-only generate, under-reporting acceptance length.
         """
-        if self.dflash_offline:
-            raise RuntimeError(
-                "DSpark offline model cannot run AR validation / pseudo_speculative_generate — "
-                "base model layers were deleted during offline conversion. Reload the full "
-                "base model before running AR validation."
-            )
-        model_output = self._base_model(input_ids=input_ids, output_hidden_states=True)
-        base_logits = self._base_model_lm_head(model_output.last_hidden_state)
-        base_token = base_logits[:, -1:, :].argmax(dim=-1).to(input_ids.device)
-
+        base_token, hidden_states = self._generation_base(input_ids)
         if steps < 1:
             return base_token, None
-
-        hid_offset = 1
-        selected = [model_output.hidden_states[lid + hid_offset] for lid in self.target_layer_ids]
-        # Same sharding hazard as HFDFlashModel.pseudo_speculative_generate: gather onto the
-        # draft's device before the cat rather than assuming one GPU holds every target layer.
-        device = self._base_device()
-        target_hidden = torch.cat([h.to(device) for h in selected], dim=-1)
-
-        block_size = self.dflash_block_size
-        bsz = input_ids.shape[0]
-
-        # Block input: anchor at position 0, mask tokens elsewhere (parallel backbone).
-        block_ids = torch.full(
-            (bsz, block_size), self.mask_token_id, dtype=torch.long, device=device
-        )
-        block_ids[:, 0] = base_token.squeeze(-1).to(device)
-        noise_embedding = self._base_model_embeddings(block_ids).to(device)
-
-        ctx_len = target_hidden.shape[1]
-        ctx_positions = torch.arange(ctx_len, device=device)
-        block_positions = torch.arange(ctx_len, ctx_len + block_size, device=device)
-        pos_ids = torch.cat([ctx_positions, block_positions]).unsqueeze(0).expand(bsz, -1)
-
-        draft_hidden = self.dflash_module(
-            noise_embedding=noise_embedding,
-            target_hidden=target_hidden,
-            position_ids=pos_ids,
-            attention_mask=self._build_generate_swa_mask(ctx_len, bsz, target_hidden.dtype, device),
-        )
+        draft_hidden = self._draft_block(base_token, hidden_states)
         backbone_logits = self._base_model_lm_head(draft_hidden)  # [B, block_size, V]
 
         # Autoregressive Markov sampling over the block.
         m = self.dflash_module
-        num_tokens = min(steps, block_size)
-        prev_token = base_token.squeeze(-1).to(device)  # anchor precedes block position 0
+        num_tokens = min(steps, self.dflash_block_size)
+        # anchor precedes block position 0
+        prev_token = base_token.squeeze(-1).to(draft_hidden.device)
         state = None
         draft_tokens = []
         for k in range(num_tokens):

From d7877eb32bf212486c7795a7f61fb62a1465f09c Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 11:43:55 +0000
Subject: [PATCH 06/15] refactor(speculative): read the DFlash family's config
 fields without a second default

Thirteen variant fields were read as getattr(config, "dflash_...", <default>) although
DFlashConfig declares every one of them with a default, so each default lived twice. One
pair had already drifted: dflash_fp32_master_weights defaults to True in the config and to
False in HFDFlashModel.modify's fallback, which a reader would take for the default. The
fallbacks never run, because modify() always receives a DFlashConfig, so the fields are now
read directly.

convert_to_dflash_model chose the registry for dflash_architecture_config.projector_type
with an if/elif chain and listed the accepted names again by hand in its error. A table
maps each projector type to its registry, and the error lists the table's keys.

No behavior change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../torch/speculative/dflash/conversion.py    | 27 ++++++++++---------
 .../torch/speculative/plugins/hf_dflash.py    |  2 +-
 .../torch/speculative/plugins/hf_dflash2.py   |  8 +++---
 .../torch/speculative/plugins/hf_domino.py    |  4 +--
 .../torch/speculative/plugins/hf_dspark.py    |  6 ++---
 .../torch/speculative/plugins/hf_lilicorr.py  |  6 ++---
 6 files changed, 27 insertions(+), 26 deletions(-)

diff --git a/modelopt/torch/speculative/dflash/conversion.py b/modelopt/torch/speculative/dflash/conversion.py
index 51f13f98c3e..859d960d0b2 100644
--- a/modelopt/torch/speculative/dflash/conversion.py
+++ b/modelopt/torch/speculative/dflash/conversion.py
@@ -48,6 +48,16 @@
 # registry so its wrapper (HFLiLiCorrModel) does not overwrite HFDFlashModel.
 LiLiCorrDMRegistry = _DMRegistryCls(prefix="LiLiCorr")
 
+# dflash_architecture_config.projector_type -> the registry holding that variant's wrapper.
+_REGISTRY_BY_PROJECTOR_TYPE = {
+    None: DFlashDMRegistry,
+    "dflash": DFlashDMRegistry,
+    "domino": DominoDMRegistry,
+    "dspark": DSparkDMRegistry,
+    "dflash2": DFlash2DMRegistry,
+    "lilicorr": LiLiCorrDMRegistry,
+}
+
 
 def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertReturnType:
     """Convert the model to a DFlash (or Domino) model as per `config`."""
@@ -59,22 +69,13 @@ def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertRe
     custom_config = config.dflash_architecture_config
     config.dflash_architecture_config = {**default_dflash_config, **custom_config}
 
-    # Route to the Domino registry when the architecture asks for the Domino head.
     projector_type = config.dflash_architecture_config.get("projector_type")
-    if projector_type == "domino":
-        registry = DominoDMRegistry
-    elif projector_type == "dspark":
-        registry = DSparkDMRegistry
-    elif projector_type == "dflash2":
-        registry = DFlash2DMRegistry
-    elif projector_type == "lilicorr":
-        registry = LiLiCorrDMRegistry
-    elif projector_type in (None, "dflash"):
-        registry = DFlashDMRegistry
-    else:
+    registry = _REGISTRY_BY_PROJECTOR_TYPE.get(projector_type)
+    if registry is None:
+        supported = [name for name in _REGISTRY_BY_PROJECTOR_TYPE if name is not None]
         raise ValueError(
             f"Unsupported dflash_architecture_config.projector_type: {projector_type!r}. "
-            "Expected 'dflash' (default), 'domino', 'dspark', 'dflash2' or 'lilicorr'."
+            f"Expected one of {supported} ('dflash' is the default)."
         )
 
     registry.register_subclass(type(model), "base_model_class")
diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 766ffaad7a7..d4d76d7e550 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -382,7 +382,7 @@ def modify(self, config):
         """Initialize DFlash draft module."""
         super().modify(config)
 
-        self.dflash_fp32_master_weights = getattr(config, "dflash_fp32_master_weights", False)
+        self.dflash_fp32_master_weights = config.dflash_fp32_master_weights
 
         base_config = self._base_llm_config
         # Use Qwen3Config (not generic PretrainedConfig) so rope_parameters is
diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py
index 5ddc95a6b48..6e4171c9210 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash2.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash2.py
@@ -72,10 +72,10 @@ def modify(self, config):
                 "candidate selector's rank/top-k)."
             )
         super().modify(config)
-        self.dflash_selector_loss_alpha = getattr(config, "dflash_selector_loss_alpha", 1.0)
-        self.dflash_lk_loss_type = getattr(config, "dflash_lk_loss_type", "ce")
-        self.dflash_lk_ce_scale = getattr(config, "dflash_lk_ce_scale", 1.0)
-        self.dflash_lk_ce_decay = getattr(config, "dflash_lk_ce_decay", 1.0)
+        self.dflash_selector_loss_alpha = config.dflash_selector_loss_alpha
+        self.dflash_lk_loss_type = config.dflash_lk_loss_type
+        self.dflash_lk_ce_scale = config.dflash_lk_ce_scale
+        self.dflash_lk_ce_decay = config.dflash_lk_ce_decay
         if self.dflash_lk_loss_type != "ce" and self.dflash_self_logit_distillation:
             raise ValueError(
                 f"dflash_lk_loss_type={self.dflash_lk_loss_type!r} needs the draft's "
diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py
index 91507974830..fc33a5155b6 100644
--- a/modelopt/torch/speculative/plugins/hf_domino.py
+++ b/modelopt/torch/speculative/plugins/hf_domino.py
@@ -129,8 +129,8 @@ def modify(self, config):
         # (DFlashConfig carries the two fields); updated each step by
         # DominoLambdaCallback. Defaults keep a single forward (e.g. unit tests)
         # well-defined without a scheduler.
-        self.dflash_lambda_base_start = getattr(config, "dflash_lambda_base_start", 1.0)
-        self.dflash_lambda_base_decay_ratio = getattr(config, "dflash_lambda_base_decay_ratio", 1.0)
+        self.dflash_lambda_base_start = config.dflash_lambda_base_start
+        self.dflash_lambda_base_decay_ratio = config.dflash_lambda_base_decay_ratio
         self._lambda_base = self.dflash_lambda_base_start
         if not getattr(self.dflash_module, "shift_label", True):
             raise NotImplementedError(
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index d000a3efa10..b29d11584ce 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -140,9 +140,9 @@ def modify(self, config):
         super().modify(config)
         # Three-term loss weights (DSpark only). Defaults follow the DeepSpec recipe
         # (L1/TVD-dominant) so a config that only sets the head still trains sensibly.
-        self.dflash_ce_loss_alpha = getattr(config, "dflash_ce_loss_alpha", 0.1)
-        self.dflash_l1_loss_alpha = getattr(config, "dflash_l1_loss_alpha", 0.9)
-        self.dflash_confidence_head_alpha = getattr(config, "dflash_confidence_head_alpha", 0.0)
+        self.dflash_ce_loss_alpha = config.dflash_ce_loss_alpha
+        self.dflash_l1_loss_alpha = config.dflash_l1_loss_alpha
+        self.dflash_confidence_head_alpha = config.dflash_confidence_head_alpha
         if self.dflash_confidence_head_alpha > 0 and not self.dflash_module.use_confidence_head:
             raise ValueError(
                 "dflash_confidence_head_alpha > 0 but the confidence head was not built; "
diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py
index be20a9d0f53..908881f8c20 100644
--- a/modelopt/torch/speculative/plugins/hf_lilicorr.py
+++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py
@@ -126,7 +126,7 @@ def modify(self, config):
             )
         super().modify(config)
 
-        weights = {name: float(getattr(config, name, -1.0)) for name in _TERM_WEIGHT_FIELDS}
+        weights = {name: float(getattr(config, name)) for name in _TERM_WEIGHT_FIELDS}
         unset = [name for name, value in weights.items() if value < 0.0]
         if unset:
             raise ValueError(
@@ -137,7 +137,7 @@ def modify(self, config):
         self.dflash_lilicorr_w_ce = weights["dflash_lilicorr_w_ce"]
         self.dflash_lilicorr_w_margin = weights["dflash_lilicorr_w_margin"]
         self.dflash_lilicorr_w_pen = weights["dflash_lilicorr_w_pen"]
-        self.dflash_lilicorr_w_cal = float(getattr(config, "dflash_lilicorr_w_cal", 0.0))
+        self.dflash_lilicorr_w_cal = float(config.dflash_lilicorr_w_cal)
         if not any(weights.values()) and self.dflash_lilicorr_w_cal == 0.0:
             raise ValueError(
                 "All LiLiCorr objective weights are 0, so the reranker would never "
@@ -145,7 +145,7 @@ def modify(self, config):
                 "least one of dflash_lilicorr_w_ce / _w_margin / _w_pen / _w_cal above 0."
             )
 
-        self.dflash_lilicorr_margin = float(getattr(config, "dflash_lilicorr_margin", -1.0))
+        self.dflash_lilicorr_margin = float(config.dflash_lilicorr_margin)
         if self.dflash_lilicorr_w_margin > 0.0 and self.dflash_lilicorr_margin <= 0.0:
             raise ValueError(
                 "dflash_lilicorr_w_margin > 0 requires a positive dflash_lilicorr_margin "

From 025e18ac1a39428eeab46a7e3bd91cebc6d15e10 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 14:06:40 +0000
Subject: [PATCH 07/15] docs(speculative): quote the shared offline
 AR-validation error in the skill tables

The shared generation step raises one message for the whole family, "An offline
DFlash-family model cannot run AR validation / pseudo_speculative_generate". The
known-failure rows in the DFlash and DSpark skill references still quoted the two
per-variant messages it replaced, so a search for the actual error found no row.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../skills/speculative-decoding/references/algorithms/dflash.md | 2 +-
 .../skills/speculative-decoding/references/algorithms/dspark.md | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md
index b0b72b7f45c..90d7f457369 100644
--- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md
+++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md
@@ -160,7 +160,7 @@ DFlash-specific:
 | --- | --- | --- |
 | `seq_len (N) must be divisible by block_size (B)` | `training.training_seq_len` not a multiple of `dflash_block_size` | Adjust either value, or pad |
 | `DFlash offline model cannot run eval/inference forward` | Offline conversion deletes base-model layers to save memory | Don't run eval on the offline model; reload the full base first |
-| `DFlash offline model cannot run AR validation / pseudo_speculative_generate` | Same cause, hit via AR validation | Keep `training.estimate_ar=false` and `training.ar_validate_steps=0` in offline runs |
+| `An offline DFlash-family model cannot run AR validation / pseudo_speculative_generate` | Same cause, hit via AR validation | Keep `training.estimate_ar=false` and `training.ar_validate_steps=0` in offline runs |
 | `dflash_swa_window_size (N) must be >= dflash_block_size (B)` | Config validation | Raise the window or lower the block size |
 | `The base model did not return hidden states required for DFlash training` | Base model's top-level forward ignores `output_hidden_states=True` | Usually a multimodal wrapper — needs a model-side fix |
 | `ERROR: DRAFT_CKPT_DIR=... contains no exported-checkpoint-* directory` | Upstream training produced no draft | Fix training; do not chase the smoke test |
diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md
index 18482ac7613..f857d193d3f 100644
--- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md
+++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md
@@ -152,7 +152,7 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i
 | `DSpark requires markov_rank > 0, got N` | Non-positive value | Set a positive rank |
 | `Unsupported markov_head_type: '...'. Expected 'vanilla', 'gated' or 'rnn'` | Typo or unsupported head | Use one of the three |
 | `dflash_confidence_head_alpha > 0 but the confidence head was not built` | Loss term enabled without the head | Set `dflash_architecture_config.use_confidence_head=true`, or set the alpha to 0 |
-| `DSpark offline model cannot run AR validation / pseudo_speculative_generate` | Offline/streaming conversion deleted base layers | Keep `estimate_ar=false` and `ar_validate_steps=0`; evaluate after export |
+| `An offline DFlash-family model cannot run AR validation / pseudo_speculative_generate` | Offline/streaming conversion deleted base layers | Keep `estimate_ar=false` and `ar_validate_steps=0`; evaluate after export |
 | Draft trains but acceptance length is poor | Draft dims left at defaults instead of matching the base | Set the GQA/FFN dims explicitly (see *Per-model adjustments*) |
 | Acceptance length capped despite clean training (streaming) | `EAGLE_CAPTURE_IDS` final layer wrong, or the vLLM aux-capture fix (vllm#46788) missing | Correct the ids; use a container with the fix |
 | Trainer init crash on a serve container | tensorboard absent | `training.report_to=none` |

From 6b63755e7effbabbaafc16ccf1052cb17d6d657d Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Tue, 27 Oct 2026 14:45:16 +0000
Subject: [PATCH 08/15] refactor(speculative): send every in-loss zero through
 the draft's outputs, DSpark's too

A step with nothing to supervise returns a zero loss that must still reach every trainable
draft parameter. How it reaches them depends on whether the draft has run. forward()'s early
returns, which come before it, read one element of each parameter. Once it has run, the loss
has to go through its outputs instead: with the draft as its own FSDP2 unit (the
MiniMax-M2.7 launcher layout), reading parameters FSDP2 has resharded puts that rank's
collectives out of step with its peers.

That second half of the rule lived in a comment on DFlash's loss and was repeated by hand
in Domino's, and earlier in this PR DSpark's in-loss branch moved to the parameter reads.
_draft_zero_loss now holds both halves: a caller whose draft has run passes its outputs as
`through`. DSpark passes its final and confidence logits, since the confidence head is not
behind final_logits. DFlash and Domino pass the one output they already went through.

Two-rank FSDP2 probe (gloo on CPU; base decoder layers and the draft each their own
fully_shard unit; rank 0 on DSpark's in-loss zero branch while rank 1 trains):

  parameters only (previous head)                  collective size mismatch, abort
  final output + one element per parameter         collective size mismatch, abort
  final output + p.sum() per parameter (main)      gloo timeout
  final output only                                mismatch: confidence head has no grad
  final + confidence outputs (this commit)         both ranks finish

That probe registers compute_markov_bias and compute_confidence_logits with
register_fsdp_forward_method. Without it no DSpark step runs in that layout, on main or
on this branch: the heads run outside the draft's forward and meet resharded DTensor
weights, as Domino's GRU head does. With the draft inside the root unit, the previous
head, main and this commit all finish. On this commit DFlash (CE and KD) also finishes
with the draft as its own unit, and Domino with the draft in the root unit.

Against the previous head, the 23-case A/B (one training step each, fp32 and bf16, the
degenerate-step cases included) gives bitwise-equal losses, gradients and metrics. The
DSpark DDP-coverage test of the zero-weight branch now also requires a gradient on both
outputs; it fails against the previous head, against either output-plus-parameter form
and against the final output alone.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 modelopt/torch/speculative/plugins/hf_dflash.py      | 12 ++++++++----
 modelopt/torch/speculative/plugins/hf_domino.py      |  2 +-
 modelopt/torch/speculative/plugins/hf_dspark.py      |  4 ++--
 .../unit/torch/speculative/plugins/test_hf_dspark.py |  4 ++++
 4 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index d4d76d7e550..8ef2aaadebd 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -874,12 +874,17 @@ def _draft_block(self, base_token, hidden_states):
             attention_mask=attn_mask,
         )
 
-    def _draft_zero_loss(self, device):
+    def _draft_zero_loss(self, device, through=()):
         """A zero loss that every trainable draft parameter feeds, for a step with nothing to train.
 
         DDP expects a gradient for each parameter on every rank; a rank returning a constant
-        loss would leave the others waiting on its reduction.
+        loss would leave the others waiting on its reduction. Once the draft has run, pass all
+        of its outputs as ``through`` (``None`` for an absent head) and the loss goes through
+        them: reading the parameters directly would put this rank's collectives out of step
+        when FSDP2 shards the draft as its own unit.
         """
+        if through:
+            return sum(out.sum() * 0.0 for out in through if out is not None)
         return sum(
             (p.reshape(-1)[0] * 0.0 for p in self.dflash_module.parameters() if p.requires_grad),
             torch.zeros((), device=device),
@@ -1002,8 +1007,7 @@ def _compute_loss(
             with torch.no_grad():
                 accuracy = _masked_top1(flat_logits, flat_targets, binary_eval_mask).item()
         else:
-            # Through the draft's output, so an FSDP unit around the draft still runs its backward.
-            loss = flat_logits.sum() * 0.0
+            loss = self._draft_zero_loss(flat_logits.device, through=(flat_logits,))
             accuracy = 0.0
 
         if return_terms:
diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py
index fc33a5155b6..d815bfcac86 100644
--- a/modelopt/torch/speculative/plugins/hf_domino.py
+++ b/modelopt/torch/speculative/plugins/hf_domino.py
@@ -231,7 +231,7 @@ def _compute_domino_loss(
                 accuracy = _masked_top1(flat_final, flat_targets, binary_eval_mask).item()
                 base_accuracy = _masked_top1(flat_base, flat_targets, binary_eval_mask).item()
         else:
-            loss = flat_final.sum() * 0.0
+            loss = self._draft_zero_loss(flat_final.device, through=(flat_final,))
             final_loss = loss
             base_loss = loss
             accuracy = 0.0
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index b29d11584ce..ed223d13b64 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -228,8 +228,8 @@ def _compute_dspark_loss(
         valid_count = flat_weights.sum() + 1e-6
 
         if valid_count <= 1.0:
-            # As forward()'s early return: DDP still reduces the confidence head's gradient.
-            loss = self._draft_zero_loss(flat_final.device)
+            # The confidence head is not behind final_logits.
+            loss = self._draft_zero_loss(flat_final.device, through=(flat_final, confidence_logits))
             metrics = {"ce_loss": 0.0, "l1_loss": 0.0, "confidence_loss": 0.0, "base_accuracy": 0.0}
             return loss, 0.0, metrics
 
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py
index 2dabbe37d57..c9000273988 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py
@@ -1061,6 +1061,8 @@ def test_loss_branch_with_zero_total_weight(self):
         final_logits, confidence_logits = model._apply_markov_head(
             hidden, backbone_logits, input_ids, anchors, n_blocks
         )
+        final_logits.retain_grad()
+        confidence_logits.retain_grad()
         loss, _, _ = model._compute_dspark_loss(
             backbone_logits,
             final_logits,
@@ -1073,3 +1075,5 @@ def test_loss_branch_with_zero_total_weight(self):
         )
         loss.backward()
         assert not self._ungraded(model), f"no gradient for {self._ungraded(model)}"
+        # Through the draft's outputs, which an FSDP2 unit around the draft needs.
+        assert final_logits.grad is not None and confidence_logits.grad is not None

From cdda93f04f6753d89e3a06e0fded7a44bcdd9231 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Tue, 27 Oct 2026 14:45:24 +0000
Subject: [PATCH 09/15] refactor(speculative): take the teacher's label
 positions by keyword

Earlier in this PR _teacher_logits started taking the positions of the labels it scores
and applying the next-token shift itself; it used to take the logit rows. Its name and the
argument's place stayed the same, so a caller written against the old meaning still runs,
one row early, and nothing says so. An out-of-tree caller that passes logit rows
positionally merges cleanly with this branch and then silently scores the wrong rows.

label_positions is now keyword-only. Such a call raises TypeError, and every call site
names what it passes. No behavior change.

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_dflash.py    |  7 ++++---
 .../torch/speculative/plugins/hf_dspark.py    |  4 +++-
 .../torch/speculative/plugins/hf_lilicorr.py  |  2 +-
 .../speculative/plugins/test_hf_dflash.py     | 20 +++++++++++--------
 4 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 8ef2aaadebd..5067efb240a 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -891,13 +891,14 @@ def _draft_zero_loss(self, device, through=()):
         )
 
     @torch.no_grad()
-    def _teacher_logits(self, base_outputs, label_positions, token_ids=None):
+    def _teacher_logits(self, base_outputs, *, label_positions, token_ids=None):
         """The base model's logits for the tokens at ``label_positions`` ([B, ...]) -> [B, ..., vocab].
 
         Next-token logits: the ones that predict the token at position ``i`` sit at ``i - 1``.
         Only those rows of the base hidden go through the final norm and lm_head, so no
         full-sequence logits are built from it; ``token_ids`` ([B, ..., k]) narrows the
-        projection to just those vocab entries.
+        projection to just those vocab entries. ``label_positions`` is keyword-only, so every
+        call says which positions it passes: handing it the logit rows reads one row early.
         """
         positions = (label_positions - 1).clamp(min=0)
         if base_outputs.logits is not None:
@@ -995,7 +996,7 @@ def _compute_loss(
 
         if valid_count > 1.0:
             if kd:
-                flat_teacher = self._teacher_logits(base_outputs, targets.positions)
+                flat_teacher = self._teacher_logits(base_outputs, label_positions=targets.positions)
                 flat_teacher = flat_teacher.reshape(flat_logits.shape)
                 target_soft = torch.softmax(flat_teacher, dim=-1)
                 draft_logsoft = torch.log_softmax(flat_logits, dim=-1)
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index ed223d13b64..1acc5048a2e 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -234,7 +234,9 @@ def _compute_dspark_loss(
             return loss, 0.0, metrics
 
         # Aligned target distribution: the base model's own prediction of each label.
-        flat_teacher = self._teacher_logits(base_outputs, targets.positions).reshape(-1, vocab)
+        flat_teacher = self._teacher_logits(
+            base_outputs, label_positions=targets.positions
+        ).reshape(-1, vocab)
 
         # Term 1: cross-entropy on the corrected (final) logits.
         ce_per_token = F.cross_entropy(flat_final, flat_targets, reduction="none")
diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py
index 908881f8c20..b38234d78eb 100644
--- a/modelopt/torch/speculative/plugins/hf_lilicorr.py
+++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py
@@ -520,7 +520,7 @@ def _candidate_target_logits(
         labels = labels[..., 1:]
         candidate_target_logits = self._teacher_logits(
             base_outputs,
-            labels,
+            label_positions=labels,
             token_ids=candidate_ids.reshape(bsz, n_blocks, num_slots, topk),
         )
         return candidate_target_logits.reshape(batch_blocks, num_slots, topk).float()
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index d100b4a8f79..7a360e9b344 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -1559,7 +1559,9 @@ def test_matches_the_full_sequence_projection(self):
         model = self._model()
         hidden = self._hidden(model)
         labels = self._labels()
-        got = model._teacher_logits(DFlashBaseModelOutput(None, base_hidden=hidden), labels)
+        got = model._teacher_logits(
+            DFlashBaseModelOutput(None, base_hidden=hidden), label_positions=labels
+        )
         want = self._rows(model._base_model_lm_head(hidden), labels - 1)
         torch.testing.assert_close(got, want)
 
@@ -1569,16 +1571,17 @@ def test_prenorm_hidden_gets_the_base_final_norm(self):
         labels = self._labels()
         outputs = DFlashBaseModelOutput(None, base_hidden=hidden, base_hidden_prenorm=True)
         want = self._rows(model._base_model_lm_head(model._base_model_norm(hidden)), labels - 1)
-        torch.testing.assert_close(model._teacher_logits(outputs, labels), want)
+        torch.testing.assert_close(model._teacher_logits(outputs, label_positions=labels), want)
 
     def test_token_ids_pick_entries_of_the_full_rows(self):
         model = self._model()
         outputs = DFlashBaseModelOutput(None, base_hidden=self._hidden(model))
         labels = self._labels()
         token_ids = torch.randint(0, model.config.vocab_size, (*labels.shape, 5))
-        rows = model._teacher_logits(outputs, labels)
+        rows = model._teacher_logits(outputs, label_positions=labels)
         torch.testing.assert_close(
-            model._teacher_logits(outputs, labels, token_ids), rows.gather(-1, token_ids)
+            model._teacher_logits(outputs, label_positions=labels, token_ids=token_ids),
+            rows.gather(-1, token_ids),
         )
 
     def test_handed_over_logits_are_gathered_not_recomputed(self):
@@ -1588,22 +1591,23 @@ def test_handed_over_logits_are_gathered_not_recomputed(self):
         labels = self._labels()
         token_ids = torch.randint(0, model.config.vocab_size, (*labels.shape, 5))
         want = self._rows(logits, labels - 1)
-        assert torch.equal(model._teacher_logits(outputs, labels), want)
+        assert torch.equal(model._teacher_logits(outputs, label_positions=labels), want)
         assert torch.equal(
-            model._teacher_logits(outputs, labels, token_ids), want.gather(-1, token_ids)
+            model._teacher_logits(outputs, label_positions=labels, token_ids=token_ids),
+            want.gather(-1, token_ids),
         )
 
     def test_label_zero_does_not_wrap_to_the_last_row(self):
         model = self._model()
         logits = torch.randn(1, SEQ_LEN, model.config.vocab_size)
         outputs = DFlashBaseModelOutput(None, logits=logits)
-        got = model._teacher_logits(outputs, torch.tensor([[0]]))
+        got = model._teacher_logits(outputs, label_positions=torch.tensor([[0]]))
         assert torch.equal(got[0, 0], logits[0, 0])
 
     def test_missing_base_distribution_raises(self):
         model = self._model()
         with pytest.raises(ValueError, match="base_model_hidden_states"):
-            model._teacher_logits(DFlashBaseModelOutput(None), self._labels())
+            model._teacher_logits(DFlashBaseModelOutput(None), label_positions=self._labels())
 
     def test_offline_dict_keeps_the_hidden_unprojected(self):
         d = {

From 21c218ce91141bb4797ab1eb2a233dda0024ca3a Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Tue, 27 Oct 2026 15:08:58 +0000
Subject: [PATCH 10/15] refactor(speculative): derive the first predicted block
 position once

Which block position is the first to predict a token after the anchor was written out in two
forms: `k + shift > 0` in _block_targets and `1 - self._label_shift` at each of the three
_block_position_weights calls. _first_predicted_position(shift) now holds it. Its
`max(0, 1 - shift)` also stays right for a shift above 1, where `1 - shift` would turn
negative and the weights' `[..., first:]` slices would quietly start from the block's end.

No behavior change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 modelopt/torch/speculative/plugins/hf_dflash.py | 9 +++++++--
 modelopt/torch/speculative/plugins/hf_domino.py | 3 ++-
 modelopt/torch/speculative/plugins/hf_dspark.py | 3 ++-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 5067efb240a..ad2308174cb 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -162,6 +162,11 @@ class _BlockTargets(NamedTuple):
     mask: torch.Tensor  # [B, N, bs] 1.0 where the position is supervised
 
 
+def _first_predicted_position(shift):
+    """The first block position whose label comes after the anchor, for a label shift."""
+    return max(0, 1 - shift)
+
+
 def _block_targets(input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, shift):
     """Labels of every block position and where they are supervised.
 
@@ -172,7 +177,7 @@ def _block_targets(input_ids, anchor_positions, block_keep_mask, loss_mask, bloc
     k = torch.arange(block_size, device=input_ids.device).view(1, 1, -1)
     mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float()
     mask = mask * inside.float()
-    mask = mask * (k + shift > 0).float()
+    mask = mask * (k >= _first_predicted_position(shift)).float()
     mask = mask * _block_gather(loss_mask, positions)
     return _BlockTargets(_block_gather(input_ids, positions), positions, mask)
 
@@ -981,7 +986,7 @@ def _compute_loss(
         weight_mask = _block_position_weights(
             targets.mask,
             self.dflash_loss_objective,
-            1 - self._label_shift,
+            _first_predicted_position(self._label_shift),
             self.dflash_loss_decay_factor,
             self.dflash_dpace_alpha,
             target_ce=lambda: (
diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py
index d815bfcac86..a8bb4002e46 100644
--- a/modelopt/torch/speculative/plugins/hf_domino.py
+++ b/modelopt/torch/speculative/plugins/hf_domino.py
@@ -73,6 +73,7 @@
     _block_position_weights,
     _block_targets,
     _block_tokens,
+    _first_predicted_position,
     _masked_top1,
 )
 from .modeling_dflash import DFlashBaseModelOutput
@@ -205,7 +206,7 @@ def _compute_domino_loss(
         weight_mask = _block_position_weights(
             targets.mask,
             "decay",
-            1 - self._label_shift,
+            _first_predicted_position(self._label_shift),
             self.dflash_loss_decay_factor,
             self.dflash_dpace_alpha,
         )
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index 1acc5048a2e..8fc67f10e7c 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -74,6 +74,7 @@
     _block_position_weights,
     _block_targets,
     _block_tokens,
+    _first_predicted_position,
     _masked_top1,
 )
 from .modeling_dflash import DFlashBaseModelOutput
@@ -216,7 +217,7 @@ def _compute_dspark_loss(
         weight_mask = _block_position_weights(
             targets.mask,
             "decay",
-            1 - self._label_shift,
+            _first_predicted_position(self._label_shift),
             self.dflash_loss_decay_factor,
             self.dflash_dpace_alpha,
         )

From fc3a81cfd65e746eec5e75e430d893da12a87e61 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Tue, 27 Oct 2026 15:10:38 +0000
Subject: [PATCH 11/15] test(speculative): DFlash's and Domino's zero-loss
 steps reach every draft parameter

Only DSpark's zero losses had a unit test. DFlash's in-loss zero, Domino's in-loss zero and
Domino's early return ran in none, and codecov reported those three lines uncovered.

The new test runs a forward with no anchorable answer, and one whose single kept block has
no label in the mask. It requires a gradient on every trainable draft parameter. Once the
draft has run, it also requires one on the draft's output, the route an FSDP2 unit around
the draft needs. Switching either in-loss zero to parameter reads fails it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../speculative/plugins/test_hf_dflash.py     | 45 +++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index 7a360e9b344..0d5ee4b1712 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -1529,6 +1529,51 @@ def test_block_targets_follow_the_label_shift(model_cls, expected):
     assert torch.equal(targets.mask, torch.tensor([mask], dtype=torch.float32))
 
 
+# DSpark's zero losses, whose confidence head adds a second output, are tested in
+# test_hf_dspark.py::TestDdpGradientCoverage.
+_ZERO_LOSS_HEADS = {
+    "dflash": {},
+    "domino": {"projector_type": "domino", "gru_hidden_dim": 32, "emb_dim": 16},
+}
+
+
+@pytest.mark.parametrize("draft_runs", [False, True], ids=["no_block", "zero_weight"])
+@pytest.mark.parametrize("head", list(_ZERO_LOSS_HEADS))
+def test_a_step_with_nothing_to_train_reaches_every_draft_parameter(head, draft_runs):
+    """DDP needs each trainable draft parameter's gradient on every rank, even at zero loss.
+
+    Once the draft has run, the loss must reach it through its output, which an FSDP2 unit
+    around the draft needs.
+    """
+    model = get_tiny_llama(num_hidden_layers=4)
+    config = get_dflash_config()
+    config["dflash_architecture_config"].update(_ZERO_LOSS_HEADS[head])
+    mtsp.convert(model, [("dflash", config)])
+    model.train()
+    outputs = []
+
+    def keep_output(module, args, output):
+        output.retain_grad()
+        outputs.append(output)
+
+    model.dflash_module.register_forward_hook(keep_output)
+    input_ids = torch.randint(1, model.config.vocab_size, (2, SEQ_LEN))
+    loss_mask = torch.zeros(2, SEQ_LEN)
+    if draft_runs:
+        # Anchors are drawn from the loss mask, so each row keeps one block, anchored at 2
+        # or 9, and none of its labels is in the mask.
+        loss_mask[:, [2, 9]] = 1.0
+    out = model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids), loss_mask=loss_mask)
+    assert out.loss.item() == 0.0
+    out.loss.backward()
+    ungraded = [
+        n for n, p in model.dflash_module.named_parameters() if p.requires_grad and p.grad is None
+    ]
+    assert not ungraded, f"no gradient for {ungraded}"
+    assert len(outputs) == int(draft_runs)
+    assert all(output.grad is not None for output in outputs)
+
+
 class TestTeacherLogits:
     """The base distribution is projected only at the rows, and vocab entries, asked for.
 

From e13d61863cc3b16d9123c9fd8b668f7b95a42710 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Tue, 27 Oct 2026 15:12:08 +0000
Subject: [PATCH 12/15] test(speculative): convert once per generation-mask
 config, not per context length

ctx_len only feeds _build_generate_swa_mask, so parametrizing over it repeated the same
tiny-model conversion for each of three context lengths. The test now converts six models
instead of eighteen and still compares the mask at the same 18 points. A failure names the
context length.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../speculative/plugins/test_hf_dflash.py      | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index 0d5ee4b1712..55d2978a7e2 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -592,22 +592,22 @@ def _legacy_generate_swa_mask(model, ctx_len, bsz, dtype, device):
 class TestGenerateMaskRule:
     """Generation evaluates the training rule for its one block, anchored at ``ctx_len``."""
 
-    @pytest.mark.parametrize("ctx_len", [1, 5, 16])
     @pytest.mark.parametrize("window", [None, BLOCK_SIZE, 2 * BLOCK_SIZE])
     @pytest.mark.parametrize("attention", ["bidirectional", "causal"])
-    def test_matches_the_hand_written_mask(self, ctx_len, window, attention):
+    def test_matches_the_hand_written_mask(self, window, attention):
         model = get_tiny_llama(num_hidden_layers=4)
         config = get_dflash_config(block_size=BLOCK_SIZE)
         config["dflash_swa_window_size"] = window
         config["dflash_draft_attention"] = attention
         mtsp.convert(model, [("dflash", config)])
-        args = (ctx_len, 2, torch.float32, torch.device("cpu"))
-        got = model._build_generate_swa_mask(*args)
-        want = _legacy_generate_swa_mask(model, *args)
-        if want is None:
-            assert got is None
-        else:
-            assert torch.equal(got, want)
+        for ctx_len in (1, 5, 16):
+            args = (ctx_len, 2, torch.float32, torch.device("cpu"))
+            got = model._build_generate_swa_mask(*args)
+            want = _legacy_generate_swa_mask(model, *args)
+            if want is None:
+                assert got is None, f"ctx_len={ctx_len}"
+            else:
+                assert torch.equal(got, want), f"ctx_len={ctx_len}"
 
 
 class TestDFlashSwaMask:

From 986fe114499d1ef128903d543e62da50eaccd657 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Wed, 28 Oct 2026 08:19:31 +0000
Subject: [PATCH 13/15] refactor(speculative): say which parameters each
 zero-loss route reaches; refuse an empty one

_draft_zero_loss's summary line promised a gradient for every trainable draft parameter.
Only the parameter route delivers that: the `through` route reaches just the parameters
behind the outputs it is handed. The docstring now says so, so a variant that adds a head
knows to pass that head's output too.

A `through` whose entries are all None used to return the int 0. It now raises. Falling back
to the parameter reads there would be wrong: once the draft has run, that is the route that
puts FSDP2's collectives out of step. No caller passes one.

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_dflash.py    | 24 ++++++++++++-------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index ad2308174cb..8ec1e94ae3c 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -880,16 +880,24 @@ def _draft_block(self, base_token, hidden_states):
         )
 
     def _draft_zero_loss(self, device, through=()):
-        """A zero loss that every trainable draft parameter feeds, for a step with nothing to train.
-
-        DDP expects a gradient for each parameter on every rank; a rank returning a constant
-        loss would leave the others waiting on its reduction. Once the draft has run, pass all
-        of its outputs as ``through`` (``None`` for an absent head) and the loss goes through
-        them: reading the parameters directly would put this rank's collectives out of step
-        when FSDP2 shards the draft as its own unit.
+        """A zero loss with a backward path to the draft, for a step with nothing to train.
+
+        DDP expects a gradient for each trainable parameter on every rank; a rank returning a
+        constant loss would leave the others waiting on its reduction. Before the draft has
+        run, the loss reads one element of every trainable draft parameter. Once it has run,
+        pass every one of its outputs as ``through`` (``None`` for an absent head): the loss
+        then reaches only the parameters behind them, so an output left out leaves its
+        parameters unreduced. Reading the parameters directly instead would put this rank's
+        collectives out of step when FSDP2 shards the draft as its own unit.
         """
         if through:
-            return sum(out.sum() * 0.0 for out in through if out is not None)
+            outputs = [out for out in through if out is not None]
+            if not outputs:
+                raise ValueError(
+                    "through= holds no draft output; the parameters cannot stand in for one "
+                    "once the draft has run."
+                )
+            return sum(out.sum() * 0.0 for out in outputs)
         return sum(
             (p.reshape(-1)[0] * 0.0 for p in self.dflash_module.parameters() if p.requires_grad),
             torch.zeros((), device=device),

From a5d6a6b032c4e3182b962bf4f42c3e08b55c7184 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Wed, 28 Oct 2026 08:20:01 +0000
Subject: [PATCH 14/15] refactor(speculative): take the position-weight floats
 by keyword; name a missing target_ce

decay_factor and dpace_alpha are adjacent floats that every caller passed positionally, and
swapped they still give a plausible weight curve. They are keyword-only now, as
label_positions is for the teacher.

D-PACE called target_ce without checking it. A caller that selects D-PACE without one now
gets a ValueError naming it, instead of "'NoneType' object is not callable". No behavior
change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 modelopt/torch/speculative/plugins/hf_dflash.py | 14 ++++++++++----
 modelopt/torch/speculative/plugins/hf_domino.py |  4 ++--
 modelopt/torch/speculative/plugins/hf_dspark.py |  4 ++--
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 8ec1e94ae3c..5d3fc293efc 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -182,14 +182,20 @@ def _block_targets(input_ids, anchor_positions, block_keep_mask, loss_mask, bloc
     return _BlockTargets(_block_gather(input_ids, positions), positions, mask)
 
 
-def _block_position_weights(mask, scheme, first, decay_factor, dpace_alpha, target_ce=None):
+def _block_position_weights(mask, scheme, first, *, decay_factor, dpace_alpha, target_ce=None):
     """``mask`` weighted per block position, by D-PACE or by exponential decay.
 
     ``first`` is the first predicted block position. ``target_ce`` returns the draft's
-    cross-entropy on its labels ([B, N, bs]); only D-PACE calls it.
+    cross-entropy on its labels ([B, N, bs]); only D-PACE calls it. The two floats are
+    keyword-only because swapping them would still give a plausible weight curve.
     """
     block_size = mask.shape[-1]
     if scheme == "dpace" and block_size > 1:
+        if target_ce is None:
+            raise ValueError(
+                "D-PACE weighs each block position by the draft's confidence on its label, "
+                "so it needs target_ce."
+            )
         with torch.no_grad():
             # Draft confidence q_i = exp(-CE) on the label; detached, as in the paper (Eq. 9).
             confidences = torch.exp(-target_ce()[..., first:].float())
@@ -995,8 +1001,8 @@ def _compute_loss(
             targets.mask,
             self.dflash_loss_objective,
             _first_predicted_position(self._label_shift),
-            self.dflash_loss_decay_factor,
-            self.dflash_dpace_alpha,
+            decay_factor=self.dflash_loss_decay_factor,
+            dpace_alpha=self.dflash_dpace_alpha,
             target_ce=lambda: (
                 loss_per_token.detach()
                 if loss_per_token is not None
diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py
index a8bb4002e46..efc22a47c46 100644
--- a/modelopt/torch/speculative/plugins/hf_domino.py
+++ b/modelopt/torch/speculative/plugins/hf_domino.py
@@ -207,8 +207,8 @@ def _compute_domino_loss(
             targets.mask,
             "decay",
             _first_predicted_position(self._label_shift),
-            self.dflash_loss_decay_factor,
-            self.dflash_dpace_alpha,
+            decay_factor=self.dflash_loss_decay_factor,
+            dpace_alpha=self.dflash_dpace_alpha,
         )
 
         flat_final = final_logits.reshape(-1, final_logits.size(-1))
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index 8fc67f10e7c..bf0c5d004f5 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -218,8 +218,8 @@ def _compute_dspark_loss(
             targets.mask,
             "decay",
             _first_predicted_position(self._label_shift),
-            self.dflash_loss_decay_factor,
-            self.dflash_dpace_alpha,
+            decay_factor=self.dflash_loss_decay_factor,
+            dpace_alpha=self.dflash_dpace_alpha,
         )
 
         flat_final = final_logits.reshape(-1, vocab)

From 112123e3156b2d60b3c85cff8581779836ca4564 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Fri, 30 Oct 2026 16:23:39 +0000
Subject: [PATCH 15/15] docs(speculative): leave the skill tables' offline
 AR-validation rows as they were

Reverts 025e18ac1a. Its two rows quote the offline AR-validation error, and HF AR
validation is being deprecated in favour of serving-engine evaluation, so the speculative
decoding skill is left for that change. The rows quote the per-variant messages again;
both still share "cannot run AR validation / pseudo_speculative_generate" with the
family-wide message, and their cause and fix still hold.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../skills/speculative-decoding/references/algorithms/dflash.md | 2 +-
 .../skills/speculative-decoding/references/algorithms/dspark.md | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md
index 90d7f457369..b0b72b7f45c 100644
--- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md
+++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md
@@ -160,7 +160,7 @@ DFlash-specific:
 | --- | --- | --- |
 | `seq_len (N) must be divisible by block_size (B)` | `training.training_seq_len` not a multiple of `dflash_block_size` | Adjust either value, or pad |
 | `DFlash offline model cannot run eval/inference forward` | Offline conversion deletes base-model layers to save memory | Don't run eval on the offline model; reload the full base first |
-| `An offline DFlash-family model cannot run AR validation / pseudo_speculative_generate` | Same cause, hit via AR validation | Keep `training.estimate_ar=false` and `training.ar_validate_steps=0` in offline runs |
+| `DFlash offline model cannot run AR validation / pseudo_speculative_generate` | Same cause, hit via AR validation | Keep `training.estimate_ar=false` and `training.ar_validate_steps=0` in offline runs |
 | `dflash_swa_window_size (N) must be >= dflash_block_size (B)` | Config validation | Raise the window or lower the block size |
 | `The base model did not return hidden states required for DFlash training` | Base model's top-level forward ignores `output_hidden_states=True` | Usually a multimodal wrapper — needs a model-side fix |
 | `ERROR: DRAFT_CKPT_DIR=... contains no exported-checkpoint-* directory` | Upstream training produced no draft | Fix training; do not chase the smoke test |
diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md
index f857d193d3f..18482ac7613 100644
--- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md
+++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md
@@ -152,7 +152,7 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i
 | `DSpark requires markov_rank > 0, got N` | Non-positive value | Set a positive rank |
 | `Unsupported markov_head_type: '...'. Expected 'vanilla', 'gated' or 'rnn'` | Typo or unsupported head | Use one of the three |
 | `dflash_confidence_head_alpha > 0 but the confidence head was not built` | Loss term enabled without the head | Set `dflash_architecture_config.use_confidence_head=true`, or set the alpha to 0 |
-| `An offline DFlash-family model cannot run AR validation / pseudo_speculative_generate` | Offline/streaming conversion deleted base layers | Keep `estimate_ar=false` and `ar_validate_steps=0`; evaluate after export |
+| `DSpark offline model cannot run AR validation / pseudo_speculative_generate` | Offline/streaming conversion deleted base layers | Keep `estimate_ar=false` and `ar_validate_steps=0`; evaluate after export |
 | Draft trains but acceptance length is poor | Draft dims left at defaults instead of matching the base | Set the GQA/FFN dims explicitly (see *Per-model adjustments*) |
 | Acceptance length capped despite clean training (streaming) | `EAGLE_CAPTURE_IDS` final layer wrong, or the vLLM aux-capture fix (vllm#46788) missing | Correct the ids; use a container with the fix |
 | Trainer init crash on a serve container | tensorboard absent | `training.report_to=none` |