From b9ec92ecf123b8e252a9671ae357224f2102a795 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 12:25:18 +0000
Subject: [PATCH 1/6] refactor(speculative): build the DFlash block objective
 from shared divergence terms

The block objective was a branch in HFDFlashModel._compute_loss (KD or CE), and DFlash2's
dflash_lk_loss_type replaced it after the fact: _lk_loss recomputed the loss from the
context's per-token cross-entropy and weights, outside the gate that turns a step with
nothing to train into a zero loss.

The divergences are now components in dflash_losses.py that name no variant, a closed
registry of callables (ctx, source) -> 0-dim loss, each reducing under the step's position
weights the way the old code did: ce, kd (soft cross-entropy against the base model), tvd
(L1 to the base model's distribution, DSpark's term), tv (1 - q(gold)) and lambda (the
acceptance-weighted CE/TV blend, still reduced term by term and blended with a detached
scalar). Each term declares whether it reads the teacher and whether it backpropagates
through the per-token cross-entropy.

An objective is an ordered [(name, weight, term)] list that modify() builds through the
_build_objective hook: DFlash and LiLiCorr keep kd or ce, DFlash2 picks tv or lambda.
_compute_loss sums it left to right and reports each term as loss_<name>, plus the
backbone's accuracy when a head sits on it. The declarations replace the per-variant
flags: offline training reconstructs the teacher's logits when a term needs them, and the
context computes the cross-entropy with gradient only when a term backpropagates through
it. LossContext gains the per-token TVD, and the TVD chunk compiles under
dflash_use_torch_compile for every variant, no longer only DSpark's modify().

One disclosed behavior change: a DFlash2 step under tv or lambda whose total position
weight is <= 1 now takes the zero-loss gate like every other objective; before, the blend
was computed over that tiny weight sum. The selector term is unchanged, still computed on
every step.

Verified against the previous commit with one train step and one generation step per case,
in fp32 and bf16: DFlash (KD and CE), DFlash2 (ce/kd/tv/lambda with the selector) and
LiLiCorr, offline and online, D-PACE and decay, including zero-loss steps, give
bitwise-equal loss, gradients, train_acc, metrics and generated tokens (30/32). The other
two runs are the change above on a step that keeps one label of weight < 1: the fp32 loss
goes from 0.96958977 to 0 (bf16: 0.968748331 to 0), the 33 backbone gradients lose the TV
term, and the 3 selector gradients and all metrics are unchanged. The new loss_<name>
metrics equal the loss for DFlash and origin_loss for LiLiCorr, bitwise.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../speculative/plugins/dflash_losses.py      | 136 +++++++++++++++++-
 .../torch/speculative/plugins/hf_dflash.py    |  64 ++++++---
 .../torch/speculative/plugins/hf_dflash2.py   |  58 +++-----
 .../torch/speculative/plugins/hf_dspark.py    |   7 -
 .../torch/speculative/plugins/hf_lilicorr.py  |   1 +
 5 files changed, 196 insertions(+), 70 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/dflash_losses.py b/modelopt/torch/speculative/plugins/dflash_losses.py
index b2e8279c3c3..9b71fb1cf42 100644
--- a/modelopt/torch/speculative/plugins/dflash_losses.py
+++ b/modelopt/torch/speculative/plugins/dflash_losses.py
@@ -41,6 +41,10 @@
 A training step's supervision: which token each block position is trained on and where
 (:func:`_block_targets`), how the positions are weighted (:func:`_block_position_weights`),
 and :class:`LossContext`, which carries them with the draft's logits to the loss.
+
+The divergences a block objective is built from live in :data:`DIVERGENCES`. An objective
+is an ordered list of ``(name, weight, term)`` entries, summed left to right by
+:func:`objective_loss`.
 """
 
 import functools
@@ -50,7 +54,14 @@
 import torch.nn.functional as F
 import torch.utils.checkpoint
 
-__all__ = ["DraftLogits", "LossContext", "LossOutput"]
+__all__ = [
+    "DIVERGENCES",
+    "DraftLogits",
+    "LossContext",
+    "LossOutput",
+    "objective_loss",
+    "objective_term",
+]
 
 
 def _block_positions(anchor_positions, block_size, offset, seq_len):
@@ -228,7 +239,7 @@ class LossContext:
     objective backpropagates through it; otherwise it runs under ``no_grad``.
     """
 
-    def __init__(self, batch, logits, targets, *, weighting, teacher, ce_grad):
+    def __init__(self, batch, logits, targets, *, weighting, teacher, ce_grad, tvd_chunk):
         """Weigh the step's block positions.
 
         Args:
@@ -238,13 +249,16 @@ def __init__(self, batch, logits, targets, *, weighting, teacher, ce_grad):
             weighting: ``_block_position_weights``' scheme, first, decay_factor and dpace_alpha.
             teacher: Returns the base model's logits for every label position.
             ce_grad: Whether the objective backpropagates through ``ce``.
+            tvd_chunk: The per-chunk kernel ``tvd`` runs, possibly compiled.
         """
         self.batch = batch
         self.logits = logits
         self.targets = targets
         self._teacher_fn = teacher
         self._ce_grad = ce_grad and torch.is_grad_enabled()
+        self._tvd_chunk = tvd_chunk
         self._ce = {}
+        self._tvd = {}
         self._teacher = None
         self.weights = _block_position_weights(
             targets.mask,
@@ -277,3 +291,121 @@ def teacher(self):
         if self._teacher is None:
             self._teacher = self._teacher_fn()
         return self._teacher
+
+    def tvd(self, source):
+        """Per-token ``||softmax(logits.<source>) - softmax(teacher)||_1``: [B * N * bs]."""
+        if source not in self._tvd:
+            logits = self.flat_logits(source)
+            self._tvd[source] = _tvd_per_token(
+                logits, self.teacher().reshape(logits.shape), chunk_fn=self._tvd_chunk
+            )
+        return self._tvd[source]
+
+
+class _Divergence:
+    """A loss term on one logits source: ``term(ctx, source)`` -> 0-dim loss.
+
+    Each term reduces under the step's position weights itself. ``needs_teacher``: it reads
+    the base model's distribution. ``needs_ce_grad``: it backpropagates through ``ctx.ce``.
+    """
+
+    needs_teacher = False
+    needs_ce_grad = False
+
+
+class _CrossEntropy(_Divergence):
+    """``-log q(gold)``: cross-entropy on the labels."""
+
+    needs_ce_grad = True
+
+    def __call__(self, ctx, source):
+        return (ctx.ce(source) * ctx.weights.view(-1)).sum() / ctx.weight_sum
+
+
+class _SoftCrossEntropy(_Divergence):
+    """Cross-entropy against the base model's distribution (logit distillation)."""
+
+    needs_teacher = True
+
+    def __call__(self, ctx, source):
+        logits = ctx.flat_logits(source)
+        target_soft = torch.softmax(ctx.teacher().reshape(logits.shape), dim=-1)
+        draft_logsoft = torch.log_softmax(logits, dim=-1)
+        per_token = -(target_soft * draft_logsoft).sum(dim=-1)
+        return (per_token * ctx.weights.view(-1)).sum() / ctx.weight_sum
+
+
+class _TeacherTotalVariation(_Divergence):
+    """``||softmax(draft) - softmax(teacher)||_1``: the L1 distance to the base model's distribution."""
+
+    needs_teacher = True
+
+    def __call__(self, ctx, source):
+        return (ctx.tvd(source) * ctx.weights.view(-1)).sum() / ctx.weight_sum
+
+
+class _LabelTotalVariation(_Divergence):
+    """``1 - q(gold)``: total variation to the one-hot label, the per-position rejection rate."""
+
+    needs_ce_grad = True
+
+    def __call__(self, ctx, source):
+        gold_probability = torch.exp(-ctx.ce(source))
+        return ((1.0 - gold_probability) * ctx.weights.view(-1)).sum() / ctx.weight_sum
+
+
+class _AcceptanceBlend(_Divergence):
+    """Cross-entropy and label total variation, blended by how often the draft is accepted.
+
+    ::
+
+        q     = exp(-ce)                 draft probability of the gold token
+        L_ce  = <ce>_w                   the cross-entropy term
+        L_tv  = <1 - q>_w                the label total-variation term
+        a     = <q>_mask                 mean acceptance over supervised positions
+        L     = s*exp(-d*a) * L_ce + (1 - s*exp(-d*a)) * L_tv
+
+    ``<.>_w`` averages under the position weights, ``<.>_mask`` under the unweighted
+    supervision mask, as the reported accuracy is. The blend weight is detached, so it
+    reshapes the objective without adding a gradient path of its own.
+    """
+
+    needs_ce_grad = True
+
+    def __init__(self, scale, decay):
+        """Blend with ``s = scale`` and ``d = decay``."""
+        self.scale = scale
+        self.decay = decay
+
+    def __call__(self, ctx, source):
+        tv_loss = _LabelTotalVariation()(ctx, source)
+        ce_loss = _CrossEntropy()(ctx, source)
+        mask = ctx.targets.mask.view(-1)
+        acceptance = (torch.exp(-ctx.ce(source)).detach() * mask).sum() / (mask.sum() + 1e-6)
+        ce_share = self.scale * torch.exp(-self.decay * acceptance)
+        return ce_share * ce_loss + (1.0 - ce_share) * tv_loss
+
+
+# The closed set of terms an objective is built from, by the name its metric reports.
+DIVERGENCES = {
+    "ce": _CrossEntropy,
+    "kd": _SoftCrossEntropy,
+    "tvd": _TeacherTotalVariation,
+    "tv": _LabelTotalVariation,
+    "lambda": _AcceptanceBlend,
+}
+
+
+def objective_term(name, weight, **params):
+    """One objective entry, ``(name, weight, term)``, with ``params`` bound to the term."""
+    return name, weight, DIVERGENCES[name](**params)
+
+
+def objective_loss(objective, ctx, source):
+    """The objective on ``logits.<source>``, and each term's value as ``loss_<name>``.
+
+    Summed left to right from 0, so a single term of weight 1.0 is that term exactly.
+    """
+    values = [(name, weight, term(ctx, source)) for name, weight, term in objective]
+    loss = sum(weight * value for _, weight, value in values)
+    return loss, {f"loss_{name}": value.detach() for name, _, value in values}
diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 93965b2b6e4..a0c2cc34bb1 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -85,7 +85,16 @@
 from ..dflash.conversion import DFlashDMRegistry
 from ..dflash.dflash_model import DFlashModel
 from .dflash_flex_attention import build_draft_block_mask
-from .dflash_losses import DraftLogits, LossContext, LossOutput, _block_targets, _masked_top1
+from .dflash_losses import (
+    DraftLogits,
+    LossContext,
+    LossOutput,
+    _block_targets,
+    _masked_top1,
+    _tvd_chunk,
+    objective_loss,
+    objective_term,
+)
 from .modeling_dflash import (  # noqa: F401
     DFlashAttention,
     DFlashBaseModelOutput,
@@ -220,7 +229,7 @@ def _base_model_norm(self):
     @property
     def _needs_base_logits(self) -> bool:
         """Whether the loss reads the target's logits; offline training then reconstructs them."""
-        return bool(self.dflash_self_logit_distillation)
+        return any(term.needs_teacher for _, _, term in self._objective)
 
     @property
     def _base_llm_config(self):
@@ -450,10 +459,22 @@ def modify(self, config):
         if self.dflash_use_torch_compile:
             self.dflash_module.compile_body()
 
+        self._objective = self._build_objective(config)
+        # Compiling fuses the TVD chunk's six vocab-wide elementwise passes.
+        self._tvd_chunk_fn = (
+            torch.compile(_tvd_chunk, dynamic=False, fullgraph=True)
+            if self.dflash_use_torch_compile
+            else _tvd_chunk
+        )
+
     def _build_draft_module(self, dflash_config):
         """Build the draft module. Subclasses override to use an augmented module."""
         return DFlashModule(dflash_config)
 
+    def _build_objective(self, config):
+        """The block objective, ``[(name, weight, term)]``: KD from the base model, or CE."""
+        return [objective_term("kd" if config.dflash_self_logit_distillation else "ce", 1.0)]
+
     def _base_device(self):
         """Device the frozen base model sits on. ``meta`` during a from_pretrained restore."""
         if self.dflash_offline:
@@ -900,33 +921,36 @@ def _loss_context(self, batch):
                 "dpace_alpha": self.dflash_dpace_alpha,
             },
             teacher=functools.partial(self._teacher_logits, batch.base_outputs, targets.positions),
-            # The KD path optimizes the soft target; its D-PACE confidences need no gradient.
-            ce_grad=not self.dflash_self_logit_distillation,
+            ce_grad=any(term.needs_ce_grad for _, _, term in self._objective),
+            tvd_chunk=self._tvd_chunk_fn,
         )
 
     def _compute_loss(self, ctx):
-        """The weighted cross-entropy (or KD) loss on the served logits, and its accuracy.
+        """The objective on the served logits and its accuracy.
 
-        A step whose position weights leave nothing to train returns a zero loss and
-        accuracy 0.0.
+        Metrics: each objective term, and the backbone's accuracy when a head sits on it.
+        A step whose position weights leave nothing to train gets ``_zero_output``.
         """
-        final = ctx.logits.final
+        final, backbone = ctx.logits
         if not ctx.trainable:
             # Through the draft's output, so an FSDP unit around the draft still runs its backward.
-            return LossOutput(final.sum() * 0.0, 0.0, {})
-
-        if self.dflash_self_logit_distillation:
-            flat_logits = ctx.flat_logits("final")
-            target_soft = torch.softmax(ctx.teacher().reshape(flat_logits.shape), dim=-1)
-            draft_logsoft = torch.log_softmax(flat_logits, dim=-1)
-            per_token = -(target_soft * draft_logsoft).sum(dim=-1)
-        else:
-            per_token = ctx.ce("final")
-        loss = (per_token * ctx.weights.view(-1)).sum() / ctx.weight_sum
-
+            return self._zero_output(ctx, final.sum() * 0.0)
+        loss, metrics = objective_loss(self._objective, ctx, "final")
         with torch.no_grad():
             accuracy = _masked_top1(final, ctx.targets.token_ids, ctx.targets.mask)
-        return LossOutput(loss, accuracy.item(), {})
+            if backbone is not final:
+                metrics["backbone_accuracy"] = _masked_top1(
+                    backbone, ctx.targets.token_ids, ctx.targets.mask
+                )
+        return LossOutput(loss, accuracy.item(), metrics)
+
+    def _zero_output(self, ctx, loss):
+        """``loss``, a zero, for a step with nothing to train; accuracy and metrics read 0."""
+        zero = torch.zeros((), device=loss.device)
+        names = [f"loss_{name}" for name, _, _ in self._objective]
+        if ctx.logits.backbone is not ctx.logits.final:
+            names.append("backbone_accuracy")
+        return LossOutput(loss, 0.0, dict.fromkeys(names, zero))
 
     def forward(
         self,
diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py
index 5d8760133e5..f716f4eaad1 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash2.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash2.py
@@ -18,7 +18,7 @@
 DFlash2 differs from DFlash only in the draft module (grouped dynamic convolutions
 around every sublayer, plus a candidate selector) and in one extra loss term, so
 this wrapper reuses ``HFDFlashModel``'s forward wholesale and overrides just
-:meth:`_compute_loss`.
+:meth:`_build_objective` and :meth:`_compute_loss`.
 
 The convolutions need no supervision of their own: they sit inside the backbone
 and are trained by the backbone loss. The selector does, because at serving time
@@ -39,7 +39,7 @@
 from transformers import PreTrainedModel
 
 from ..dflash.conversion import DFlash2DMRegistry
-from .dflash_losses import LossOutput, _block_tokens
+from .dflash_losses import LossOutput, _block_tokens, objective_term
 from .hf_dflash import HFDFlashModel
 from .modeling_dflash2 import DFlash2Module
 
@@ -74,12 +74,9 @@ def modify(self, config):
             )
         super().modify(config)
         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:
+        if config.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 "
+                f"dflash_lk_loss_type={config.dflash_lk_loss_type!r} needs the draft's "
                 "probability of the gold token, which the KD path never forms -- it "
                 "optimizes a soft target instead. Set dflash_self_logit_distillation=false, "
                 "or dflash_lk_loss_type='ce' to keep distillation."
@@ -148,35 +145,17 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas
         # every training step. The trainer converts them at the logging boundary.
         return loss, accuracy.detach(), coverage.detach()
 
-    def _lk_loss(self, ctx):
-        """Re-weight the block objective between cross-entropy and acceptance.
-
-        Both terms are read off the same per-position cross-entropy the backbone loss
-        already produced, so the target alignment and position weighting are shared::
-
-            q     = exp(-ce)                 draft probability of the gold token
-            L_ce  = <ce>_w                   today's objective
-            L_tv  = <1 - q>_w                total variation to the one-hot target,
-                                             i.e. the per-position acceptance loss
-            a     = <q>_mask                 mean acceptance over supervised positions
-            L     = s*exp(-d*a) * L_ce + (1 - s*exp(-d*a)) * L_tv
-
-        ``<.>_w`` averages under the position weighting, ``<.>_mask`` under the
-        unweighted supervised mask, matching how the reported accuracy is normalized.
-        The blend weight is detached, so it reshapes the objective without adding a
-        gradient path of its own.
-        """
-        ce, weights, weight_sum = ctx.ce("final"), ctx.weights.view(-1), ctx.weight_sum
-        gold_probability = torch.exp(-ce)
-        tv_loss = ((1.0 - gold_probability) * weights).sum() / weight_sum
-        if self.dflash_lk_loss_type == "tv":
-            return tv_loss
-
-        ce_loss = (ce * weights).sum() / weight_sum
-        mask = ctx.targets.mask.view(-1)
-        acceptance = (gold_probability.detach() * mask).sum() / (mask.sum() + 1e-6)
-        ce_share = self.dflash_lk_ce_scale * torch.exp(-self.dflash_lk_ce_decay * acceptance)
-        return ce_share * ce_loss + (1.0 - ce_share) * tv_loss
+    def _build_objective(self, config):
+        """DFlash's objective, or the label TV / CE-TV blend ``dflash_lk_loss_type`` selects."""
+        if config.dflash_lk_loss_type == "tv":
+            return [objective_term("tv", 1.0)]
+        if config.dflash_lk_loss_type == "lambda":
+            return [
+                objective_term(
+                    "lambda", 1.0, scale=config.dflash_lk_ce_scale, decay=config.dflash_lk_ce_decay
+                )
+            ]
+        return super()._build_objective(config)
 
     def _compute_loss(self, ctx):
         """Backbone DFlash loss plus the candidate-selector cross-entropy.
@@ -185,11 +164,8 @@ def _compute_loss(self, ctx):
         comparable; the selector's own accuracy and coverage are added to the metrics.
         """
         out = super()._compute_loss(ctx)
-        loss = out.loss
-        if self.dflash_lk_loss_type != "ce":
-            loss = self._lk_loss(ctx)
         if self.dflash_selector_loss_alpha <= 0:
-            return out._replace(loss=loss)
+            return out
 
         batch, targets = ctx.batch, ctx.targets
         bsz, n_blocks = batch.anchors.shape
@@ -223,5 +199,5 @@ def _compute_loss(self, ctx):
             "selector_coverage": selector_coverage,
         }
         return LossOutput(
-            loss + self.dflash_selector_loss_alpha * selector_loss, out.accuracy, metrics
+            out.loss + self.dflash_selector_loss_alpha * selector_loss, out.accuracy, metrics
         )
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index 0891f9a924b..a65897b8d17 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -73,7 +73,6 @@
     _block_targets,
     _block_tokens,
     _masked_top1,
-    _tvd_chunk,
     _tvd_per_token,
 )
 from .hf_dflash import HFDFlashModel
@@ -118,12 +117,6 @@ def modify(self, config):
                 "dflash_confidence_head_alpha > 0 but the confidence head was not built; "
                 "set dflash_architecture_config.use_confidence_head=true."
             )
-        # Compiling fuses the chunk's six vocab-wide elementwise passes.
-        self._tvd_chunk_fn = (
-            torch.compile(_tvd_chunk, dynamic=False, fullgraph=True)
-            if self.dflash_use_torch_compile
-            else _tvd_chunk
-        )
 
     def get_exporter(self):
         """Get the exporter for the DSpark draft model."""
diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py
index 97c8f6aea25..be3100e0278 100644
--- a/modelopt/torch/speculative/plugins/hf_lilicorr.py
+++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py
@@ -214,6 +214,7 @@ def _compute_loss(self, ctx):
         )
         # The identity `loss == origin_loss + lilicorr_loss` is the cheap parity check
         # on this objective, so both halves are reported next to the total.
+        metrics.update(out.metrics)
         metrics["origin_loss"] = out.loss.detach().float()
         metrics["origin_accuracy"] = torch.full((), out.accuracy, device=out.loss.device)
         return LossOutput(out.loss + lilicorr_loss, accuracy, metrics)

From 7fc13afb6ba5bfa628e21b8711c6e658ea83b6ce Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 12:27:23 +0000
Subject: [PATCH 2/6] refactor(speculative): train DSpark through the shared
 DFlash forward

HFDSparkModel.forward was a copy of the DFlash training forward (base model, loss mask,
anchors, draft inputs, draft) followed by _apply_markov_head and _compute_dspark_loss,
which rebuilt the labels, weights, gate and accuracy that the shared loss now provides.
Copies of the forward are where the family's fixes went missing: the loss-mask rule DFlash
gained for padding never reached it.

DSpark now changes only its hooks:
- _build_objective: [("ce", dflash_ce_loss_alpha), ("tvd", dflash_l1_loss_alpha)].
- _draft_logits: the Markov transition bias on the backbone's logits (the old
  _apply_markov_head without the confidence head).
- _online_base_outputs: the inner model and its final hidden, which the loss projects
  only at the rows it reads.
- _position_weighting: "decay", as before (dflash_loss_objective is not wired for DSpark).
- _compute_loss: the objective, plus the confidence head's BCE against 1 - 0.5 * TVD,
  sharing the TVD term's per-token values. A step with nothing to train keeps returning
  _draft_zero_loss, which also feeds the confidence head for DDP.

With the shared forward DSpark adopts DFlash's loss mask (labels, attention_mask and
loss_mask intersected), the fix the padding PR makes on main; it is identical whenever
attention_mask is all ones where labels are valid, as in every shipped recipe and in
streaming and offline batches. Offline batches now fail at the entry when the TVD term has
no teacher to read, rather than after the draft has run.

Metrics move to draft_metrics: ce_loss, l1_loss and base_accuracy become loss_ce, loss_tvd
and backbone_accuracy; confidence_loss keeps its name.

Verified against the previous commit with one train step and one generation step per case,
in fp32 and bf16: the vanilla and rnn heads with the confidence head on, offline and
online, and a step that takes the zero-loss branch give bitwise-equal loss, gradients,
train_acc, metrics (under the names above) and generated tokens (8/8).

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_dspark.py    | 303 ++++--------------
 1 file changed, 59 insertions(+), 244 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index a65897b8d17..69b5fdc3f1f 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -64,17 +64,9 @@
 import torch
 import torch.nn.functional as F
 from transformers import PreTrainedModel
-from transformers.trainer_pt_utils import LabelSmoother
-from transformers.utils import ModelOutput
 
 from ..dflash.conversion import DSparkDMRegistry
-from .dflash_losses import (
-    _block_position_weights,
-    _block_targets,
-    _block_tokens,
-    _masked_top1,
-    _tvd_per_token,
-)
+from .dflash_losses import DraftLogits, LossOutput, _block_tokens, objective_term
 from .hf_dflash import HFDFlashModel
 from .modeling_dflash import DFlashBaseModelOutput
 from .modeling_dspark import DSparkModule
@@ -99,7 +91,7 @@ def _build_draft_module(self, dflash_config):
         return DSparkModule(dflash_config)
 
     def modify(self, config):
-        """Initialize the DSpark draft module and read the loss-mixing weights."""
+        """Initialize the DSpark draft module and read the confidence term's weight."""
         arch_config = config.dflash_architecture_config
         if arch_config.get("markov_rank") is None:
             raise ValueError(
@@ -107,10 +99,6 @@ def modify(self, config):
                 "dflash_architecture_config (the Markov head's low-rank dimension)."
             )
         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 = 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(
@@ -124,248 +112,75 @@ def get_exporter(self):
 
         return DSparkExporter(self)
 
-    def _apply_markov_head(self, hidden, backbone_logits, input_ids, anchor_positions, n_blocks):
-        """Add the Markov transition bias to the backbone base logits.
+    def _build_objective(self, config):
+        """CE plus TVD to the base model; the DeepSpec defaults are TVD-dominant."""
+        return [
+            objective_term("ce", config.dflash_ce_loss_alpha),
+            objective_term("tvd", config.dflash_l1_loss_alpha),
+        ]
 
-        Returns ``(final_logits [B, N, bs, V], confidence_logits [B, N, bs] | None)``.
-        """
-        bsz = input_ids.shape[0]
-        bs = self.dflash_block_size
-
-        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_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
-
-        confidence_logits = None
-        if self.dflash_module.use_confidence_head:
-            confidence_logits = self.dflash_module.compute_confidence_logits(prev_ids, hidden4d)
-        return final4d, confidence_logits
-
-    def _compute_dspark_loss(
-        self,
-        backbone_logits,
-        final_logits,
-        confidence_logits,
-        input_ids,
-        anchor_positions,
-        block_keep_mask,
-        loss_mask,
-        base_outputs,
-    ):
-        """Compute the three-term DSpark loss (CE + TVD + confidence BCE) and metrics.
-
-        Uses next-token (shift_label) alignment: block position k predicts the token
-        at anchor+k+1; the aligned target distribution is the base model's own
-        next-token distribution at position anchor+k (= label index - 1), read from
-        ``base_outputs``.
-        """
-        vocab = final_logits.size(-1)
-        targets = _block_targets(
-            input_ids,
-            anchor_positions,
-            block_keep_mask,
-            loss_mask,
-            self.dflash_block_size,
-            self._label_shift,
-        )
-        binary_eval_mask = targets.mask.view(-1)
+    @property
+    def _position_weighting(self):
         # 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,
-        )
-
-        flat_final = final_logits.reshape(-1, vocab)
-        flat_base = backbone_logits.reshape(-1, vocab)
-        flat_targets = targets.token_ids.reshape(-1)
-        flat_weights = weight_mask.reshape(-1)
-        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)
-            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: the base model's own prediction of each label.
-        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")
-        ce_loss = (ce_per_token * flat_weights).sum() / valid_count
-
-        # Term 2: total-variation distance between the corrected draft and target.
-        # Chunked + checkpointed to avoid materializing two [N, vocab] softmaxes at once.
-        l1_per_token = _tvd_per_token(
-            flat_final,
-            flat_teacher,
-            chunk_fn=self._tvd_chunk_fn,
-        )
-        l1_loss = (l1_per_token * flat_weights).sum() / valid_count
-
-        # Term 3: confidence head BCE against the analytical accept rate c* = 1 - 0.5*TVD.
-        confidence_loss = ce_loss.new_zeros(())
-        if confidence_logits is not None and self.dflash_confidence_head_alpha > 0:
-            accept_rate = (1.0 - 0.5 * l1_per_token).clamp(0.0, 1.0).detach()
-            conf_bce = F.binary_cross_entropy_with_logits(
-                confidence_logits.reshape(-1).float(), accept_rate, reduction="none"
-            )
-            confidence_loss = (conf_bce * flat_weights).sum() / valid_count
-
-        loss = (
-            self.dflash_ce_loss_alpha * ce_loss
-            + self.dflash_l1_loss_alpha * l1_loss
-            + self.dflash_confidence_head_alpha * confidence_loss
-        )
+        return "decay"
 
+    def _online_base_outputs(self, input_ids, attention_mask, **kwargs):
+        """The target hidden states, and the base's final hidden for the loss to project."""
+        # The inner model, not the CausalLM: the loss projects only the rows it reads.
         with torch.no_grad():
-            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()]
-            ).tolist()
-            metrics = {
-                "ce_loss": ce_v,
-                "l1_loss": l1_v,
-                "confidence_loss": conf_v,
-                "base_accuracy": base_acc_v,
-            }
-        return loss, acc_v, metrics
-
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        labels=None,
-        use_cache=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        cache_position=None,
-        **kwargs,
-    ):
-        """DSpark training forward: DFlash backbone + Markov head + three-term loss.
-
-        Mirrors ``HFDFlashModel.forward`` for data preparation (reusing the inherited
-        anchor/noise/mask/position helpers), then applies the Markov head and the
-        DSpark loss. Eval/offline-eval is delegated to the DFlash parent.
-        """
-        if not self.training:
-            return super().forward(
+            base_out = self._base_model(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_values=past_key_values,
-                inputs_embeds=inputs_embeds,
-                use_cache=use_cache,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                cache_position=cache_position,
-                **kwargs,
-            )
-
-        bsz, seq_len = input_ids.shape
-        block_size = self.dflash_block_size
-        device = input_ids.device
-
-        if seq_len % block_size != 0:
-            raise ValueError(
-                f"seq_len ({seq_len}) must be divisible by block_size ({block_size}). "
-                f"Adjust training_seq_len or use padding."
-            )
-
-        # 1. Target hidden states, plus what the TVD/confidence terms read the base
-        #    distribution from (the loss projects only the rows it uses).
-        if self.dflash_offline:
-            assert "base_model_outputs" in kwargs
-            base_outputs = DFlashBaseModelOutput.from_offline_dict(kwargs["base_model_outputs"])
-        else:
-            # Call the inner base model directly (NOT super().forward(), which during
-            # training runs the full DFlash pipeline).
-            with torch.no_grad():
-                base_out = self._base_model(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    output_hidden_states=True,
-                )
-            base_outputs = DFlashBaseModelOutput(
-                target_hidden=self._target_hidden(base_out.hidden_states),
-                base_hidden=base_out.last_hidden_state,
+                output_hidden_states=True,
             )
-        target_hidden = base_outputs.target_hidden
-
-        # 2. Build loss mask (same convention as DFlash/Domino).
-        if labels is not None:
-            loss_mask = (labels != LabelSmoother.ignore_index).float()
-        elif attention_mask is not None:
-            loss_mask = attention_mask.float()
-        else:
-            loss_mask = torch.ones(bsz, seq_len, device=device)
-        if kwargs.get("loss_mask") is not None:
-            loss_mask = loss_mask * kwargs["loss_mask"]
-
-        # 3. Random anchor sampling (inherited).
-        anchor_positions, block_keep_mask = self._sample_anchor_positions(
-            seq_len, loss_mask, device
-        )
-        n_blocks = anchor_positions.shape[1]
-
-        if n_blocks == 0 or not block_keep_mask.any():
-            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(
-            input_ids, anchor_positions, block_keep_mask, n_blocks
-        )
-        full_pos = self._build_position_ids(seq_len, anchor_positions, device)
-        attn_mask = self._build_draft_attention_mask(
-            seq_len,
-            anchor_positions,
-            block_keep_mask,
-            n_blocks,
-            target_hidden.dtype,
-            device,
-            window=self.dflash_swa_window_size,
+        return DFlashBaseModelOutput(
+            target_hidden=self._target_hidden(base_out.hidden_states),
+            base_hidden=base_out.last_hidden_state,
         )
 
-        # 5. Draft backbone forward.
-        hidden = self.dflash_module(
-            noise_embedding=noise_embedding,
-            target_hidden=target_hidden,
-            position_ids=full_pos,
-            attention_mask=attn_mask,
-        )
+    def _head_inputs(self, batch):
+        """Each block position's teacher-forced previous token, and the draft's hidden states.
 
-        # 6. Backbone logits → Markov correction → three-term loss.
-        backbone_logits = self._base_model_lm_head(hidden).reshape(bsz, n_blocks, block_size, -1)
-        final_logits, confidence_logits = self._apply_markov_head(
-            hidden, backbone_logits, input_ids, anchor_positions, n_blocks
-        )
-        loss, accuracy, metrics = self._compute_dspark_loss(
-            backbone_logits,
-            final_logits,
-            confidence_logits,
-            input_ids,
-            anchor_positions,
-            block_keep_mask,
-            loss_mask,
-            base_outputs,
+        The previous token of position k is the real token at anchor+k (the anchor itself
+        for position 0). Returns ``prev_ids [B, N, bs]`` and ``hidden [B, N, bs, H]``.
+        """
+        bsz, n_blocks = batch.anchors.shape
+        bs = self.dflash_block_size
+        prev_ids = _block_tokens(batch.input_ids, batch.anchors, bs, self._label_shift - 1)
+        return prev_ids, batch.hidden.reshape(bsz, n_blocks, bs, -1)
+
+    def _draft_logits(self, batch):
+        """The backbone's logits plus the Markov transition bias."""
+        backbone = super()._draft_logits(batch).backbone
+        prev_ids, hidden = self._head_inputs(batch)
+        bias = self.dflash_module.compute_markov_bias(prev_ids, hidden)
+        return DraftLogits(final=backbone + bias, backbone=backbone)
+
+    def _compute_loss(self, ctx):
+        """The CE + TVD objective plus the confidence head's BCE against the accept rate."""
+        if not ctx.trainable:
+            # As forward()'s early return: DDP still reduces the confidence head's gradient.
+            out = self._zero_output(ctx, self._draft_zero_loss(ctx.logits.final.device))
+            zero = torch.zeros((), device=out.loss.device)
+            return out._replace(metrics={**out.metrics, "confidence_loss": zero})
+
+        out = super()._compute_loss(ctx)
+        confidence_loss = out.loss.new_zeros(())
+        if self.dflash_confidence_head_alpha > 0:
+            prev_ids, hidden = self._head_inputs(ctx.batch)
+            confidence_logits = self.dflash_module.compute_confidence_logits(prev_ids, hidden)
+            # The analytical accept rate c* = 1 - 0.5 * TVD.
+            accept_rate = (1.0 - 0.5 * ctx.tvd("final")).clamp(0.0, 1.0).detach()
+            conf_bce = F.binary_cross_entropy_with_logits(
+                confidence_logits.reshape(-1).float(), accept_rate, reduction="none"
+            )
+            confidence_loss = (conf_bce * ctx.weights.view(-1)).sum() / ctx.weight_sum
+        return LossOutput(
+            out.loss + self.dflash_confidence_head_alpha * confidence_loss,
+            out.accuracy,
+            {**out.metrics, "confidence_loss": confidence_loss.detach()},
         )
 
-        return ModelOutput(loss=loss, logits=None, train_acc=[[accuracy]], dspark_metrics=metrics)
-
     @torch.no_grad()
     def pseudo_speculative_generate(self, input_ids, steps=1):
         """Generate draft tokens for AR validation, with the Markov correction applied.

From 370ac96b4ce547c9bcfb973f36290403432a1210 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 12:30:14 +0000
Subject: [PATCH 3/6] refactor(speculative): train Domino through the shared
 DFlash forward

HFDominoModel.forward was the second copy of the DFlash training forward, followed by
_apply_domino_head and _compute_domino_loss, which again rebuilt the labels, weights, gate
and accuracies. Domino now changes only its hooks:
- _build_objective: [("ce", 1.0)].
- _draft_logits: the GRU correction on the backbone's logits (the old _apply_domino_head).
- _online_base_outputs: the inner model; Domino reads no teacher.
- _position_weighting: "decay", as before.
- _compute_loss: (1 - lambda_base) * objective("final") + lambda_base *
  objective("backbone"), behind the shared zero-loss gate.

Domino no longer overrides forward(): its eval branch already delegated to the shared
forward, so with training on the hooks above nothing is left to override.

Like DSpark, Domino adopts DFlash's loss mask (labels, attention_mask and loss_mask
intersected): identical whenever attention_mask is all ones where labels are valid.

Metrics move to draft_metrics: final_loss, base_loss and base_accuracy become loss_final,
loss_backbone and backbone_accuracy; lambda_base keeps its name.

Verified against the previous commit with one train step and one generation step per case,
in fp32 and bf16: offline and online steps and a step that takes the zero-loss branch give
bitwise-equal loss, gradients, train_acc, metrics (under the names above) and generated
tokens (6/6).

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_domino.py    | 233 +++---------------
 1 file changed, 41 insertions(+), 192 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py
index 2e924b121af..0a4bc9fdbd4 100644
--- a/modelopt/torch/speculative/plugins/hf_domino.py
+++ b/modelopt/torch/speculative/plugins/hf_domino.py
@@ -61,14 +61,11 @@
 import logging
 
 import torch
-import torch.nn.functional as F
 from transformers import PreTrainedModel
 from transformers.trainer_callback import TrainerCallback
-from transformers.trainer_pt_utils import LabelSmoother
-from transformers.utils import ModelOutput
 
 from ..dflash.conversion import DominoDMRegistry
-from .dflash_losses import _block_position_weights, _block_targets, _block_tokens, _masked_top1
+from .dflash_losses import DraftLogits, LossOutput, _block_tokens, objective_loss, objective_term
 from .hf_dflash import HFDFlashModel
 from .modeling_dflash import DFlashBaseModelOutput
 from .modeling_domino import DominoModule
@@ -137,28 +134,35 @@ def get_exporter(self):
     def _current_lambda_base(self) -> float:
         return float(getattr(self, "_lambda_base", self.dflash_lambda_base_start))
 
-    def _apply_domino_head(self, hidden, base_logits, input_ids, anchor_positions, n_blocks):
-        """Add the GRU causal correction to the suffix positions of each block.
+    def _build_objective(self, config):
+        """CE on the labels, for the corrected logits and for the backbone's."""
+        return [objective_term("ce", 1.0)]
 
-        Args:
-            hidden: Draft backbone output [B, N*block_size, H].
-            base_logits: Backbone logits [B, N*block_size, vocab].
-            input_ids: Original token IDs [B, seq_len].
-            anchor_positions: Anchor positions per block [B, N].
-            n_blocks: Number of blocks N.
+    @property
+    def _position_weighting(self):
+        # dflash_loss_objective is not wired for Domino yet: it always decays.
+        return "decay"
 
-        Returns:
-            Corrected logits [B, N*block_size, vocab].
-        """
-        bsz = input_ids.shape[0]
-        bs = self.dflash_block_size
+    def _online_base_outputs(self, input_ids, attention_mask, **kwargs):
+        """The target hidden states; Domino does not use target-logit KD."""
+        with torch.no_grad():
+            base_out = self._base_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+        return DFlashBaseModelOutput(target_hidden=self._target_hidden(base_out.hidden_states))
+
+    def _draft_logits(self, batch):
+        """The backbone's logits with the GRU causal correction added to each block's suffix."""
+        backbone = super()._draft_logits(batch).backbone
+        bsz, n_blocks, bs, _ = backbone.shape
         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)
+        hidden4d = batch.hidden.reshape(bsz, n_blocks, bs, -1)
 
         # Teacher-forced previous tokens: the real token at anchor+j for j in [0, bs).
-        prev_ids = _block_tokens(input_ids, anchor_positions, bs, self._label_shift - 1)
+        prev_ids = _block_tokens(batch.input_ids, batch.anchors, 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))
@@ -170,183 +174,28 @@ def _apply_domino_head(self, hidden, base_logits, input_ids, anchor_positions, n
         z_n = hidden4d[:, :, suffix_start:, :]
         logits_e = self.dflash_module.embed_proj(torch.cat([z_n, prefix_states], dim=-1))
 
-        prefix_logits = base4d[:, :, :suffix_start, :]
-        suffix_logits = base4d[:, :, suffix_start:, :] + logits_e
-        final4d = torch.cat([prefix_logits, suffix_logits], dim=2)
-        return final4d.reshape(bsz, n_blocks * bs, -1)
-
-    def _compute_domino_loss(
-        self, base_logits, final_logits, input_ids, anchor_positions, block_keep_mask, loss_mask
-    ):
-        """Compute the (1-lambda)*final + lambda*base weighted CE loss and accuracies.
-
-        Uses next-token (shift_label) alignment: position k predicts the token at
-        anchor+k+1, and position 0 is *not* excluded (unlike base DFlash).
-        """
-        targets = _block_targets(
-            input_ids,
-            anchor_positions,
-            block_keep_mask,
-            loss_mask,
-            self.dflash_block_size,
-            self._label_shift,
-        )
-        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,
+        prefix_logits = backbone[:, :, :suffix_start, :]
+        suffix_logits = backbone[:, :, suffix_start:, :] + logits_e
+        return DraftLogits(
+            final=torch.cat([prefix_logits, suffix_logits], dim=2), backbone=backbone
         )
 
-        flat_final = final_logits.reshape(-1, final_logits.size(-1))
-        flat_base = base_logits.reshape(-1, base_logits.size(-1))
-        flat_targets = targets.token_ids.reshape(-1)
-        flat_weights = weight_mask.reshape(-1)
-        valid_count = flat_weights.sum() + 1e-6
-
+    def _compute_loss(self, ctx):
+        """``(1 - lambda_base) * objective(final) + lambda_base * objective(backbone)``."""
+        out = super()._compute_loss(ctx)
         lambda_base = self._current_lambda_base()
-
-        if valid_count > 1.0:
-            final_loss = (
-                F.cross_entropy(flat_final, flat_targets, reduction="none") * flat_weights
-            ).sum() / valid_count
-            base_loss = (
-                F.cross_entropy(flat_base, flat_targets, reduction="none") * flat_weights
-            ).sum() / valid_count
-            loss = (1.0 - lambda_base) * final_loss + lambda_base * base_loss
-
-            with torch.no_grad():
-                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
-            base_loss = loss
-            accuracy = 0.0
-            base_accuracy = 0.0
-
+        loss = backbone_loss = out.loss
+        if ctx.trainable:
+            backbone_loss, _ = objective_loss(self._objective, ctx, "backbone")
+            loss = (1.0 - lambda_base) * out.loss + lambda_base * backbone_loss
         metrics = {
-            "final_loss": final_loss.detach().item(),
-            "base_loss": base_loss.detach().item(),
-            "base_accuracy": base_accuracy,
-            "lambda_base": lambda_base,
+            "loss_final": out.loss.detach(),
+            "loss_backbone": backbone_loss.detach(),
+            "backbone_accuracy": out.metrics["backbone_accuracy"],
+            # float64 holds the scheduled value exactly.
+            "lambda_base": torch.full((), lambda_base, dtype=torch.float64, device=out.loss.device),
         }
-        return loss, accuracy, metrics
-
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        labels=None,
-        use_cache=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        cache_position=None,
-        **kwargs,
-    ):
-        """Domino training forward: DFlash backbone + causal correction head + dual loss.
-
-        Mirrors ``HFDFlashModel.forward`` for data preparation (reusing the inherited
-        anchor/noise/mask/position helpers), then applies the Domino head and the
-        two-term loss. Eval/offline-eval is delegated to the DFlash parent.
-        """
-        if not self.training:
-            return super().forward(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_values=past_key_values,
-                inputs_embeds=inputs_embeds,
-                labels=labels,
-                use_cache=use_cache,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                cache_position=cache_position,
-                **kwargs,
-            )
-
-        bsz, seq_len = input_ids.shape
-        block_size = self.dflash_block_size
-        device = input_ids.device
-
-        if seq_len % block_size != 0:
-            raise ValueError(
-                f"seq_len ({seq_len}) must be divisible by block_size ({block_size}). "
-                f"Adjust training_seq_len or use padding."
-            )
-
-        # 1. Target hidden states (Domino does not use target-logit KD).
-        if self.dflash_offline:
-            assert "base_model_outputs" in kwargs
-            base_outputs = DFlashBaseModelOutput.from_offline_dict(kwargs["base_model_outputs"])
-            target_hidden = base_outputs.target_hidden
-        else:
-            with torch.no_grad():
-                base_out = self._base_model(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    output_hidden_states=True,
-                )
-            target_hidden = self._target_hidden(base_out.hidden_states)
-
-        # 2. Build loss mask (same convention as DFlash).
-        if labels is not None:
-            loss_mask = (labels != LabelSmoother.ignore_index).float()
-        elif attention_mask is not None:
-            loss_mask = attention_mask.float()
-        else:
-            loss_mask = torch.ones(bsz, seq_len, device=device)
-        if kwargs.get("loss_mask") is not None:
-            loss_mask = loss_mask * kwargs["loss_mask"]
-
-        # 3. Random anchor sampling (inherited).
-        anchor_positions, block_keep_mask = self._sample_anchor_positions(
-            seq_len, loss_mask, device
-        )
-        n_blocks = anchor_positions.shape[1]
-
-        if n_blocks == 0 or not block_keep_mask.any():
-            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(
-            input_ids, anchor_positions, block_keep_mask, n_blocks
-        )
-        full_pos = self._build_position_ids(seq_len, anchor_positions, device)
-        attn_mask = self._build_draft_attention_mask(
-            seq_len,
-            anchor_positions,
-            block_keep_mask,
-            n_blocks,
-            target_hidden.dtype,
-            device,
-            window=self.dflash_swa_window_size,
-        )
-
-        # 5. Draft backbone forward.
-        hidden = self.dflash_module(
-            noise_embedding=noise_embedding,
-            target_hidden=target_hidden,
-            position_ids=full_pos,
-            attention_mask=attn_mask,
-        )
-
-        # 6. Base + corrected logits, then dual loss.
-        base_logits = self._base_model_lm_head(hidden)
-        final_logits = self._apply_domino_head(
-            hidden, base_logits, input_ids, anchor_positions, n_blocks
-        )
-        loss, accuracy, metrics = self._compute_domino_loss(
-            base_logits, final_logits, input_ids, anchor_positions, block_keep_mask, loss_mask
-        )
-
-        return ModelOutput(loss=loss, logits=None, train_acc=[[accuracy]], domino_metrics=metrics)
+        return LossOutput(loss, out.accuracy, metrics)
 
     @torch.no_grad()
     def pseudo_speculative_generate(self, input_ids, steps=1):

From 6ca2335a410130ef4ddbde40f3dbc4c68fef767a Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Tue, 13 Oct 2026 14:18:23 +0000
Subject: [PATCH 4/6] test(speculative): follow the DFlash family onto the
 shared forward and loss

Update the call sites of the internals the previous commits removed or renamed, keeping
each test's intent:
- DSpark's zero-weight DDP test runs the real _loss_context/_compute_loss instead of the
  removed _apply_markov_head/_compute_dspark_loss.
- Domino's AR-validation test teacher-forces the drafts through _draft_logits instead of
  the removed _apply_domino_head. It still fails when either the drafting or the training
  side skips the correction head or feeds it the wrong prefix.
- The DSpark and Domino metrics are read from draft_metrics, under the new names where
  they changed, and as tensors.

One new test pins the disclosed behavior change: a DFlash2 step whose position weights
sum to <= 1 takes the zero loss under ce, tv and lambda alike. Before, tv and lambda gave
0.968 and 3.352 on it.

tests/unit/torch/speculative and test_hf_spec_rope_export.py: 516 passed on transformers
5.5.4, 5.12.1 and 5.18.0. The changed tests take at most 0.7 s each on 4 CPUs.

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_dflash2.py    | 28 +++++++++++++++
 .../speculative/plugins/test_hf_domino.py     | 29 ++++++++--------
 .../speculative/plugins/test_hf_dspark.py     | 34 +++++++------------
 3 files changed, 55 insertions(+), 36 deletions(-)

diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
index 21ca9ca59e5..bff46e08f3f 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
@@ -541,6 +541,34 @@ def test_acceptance_term_is_a_probability(self):
         loss = self._loss(dflash_lk_loss_type="tv", dflash_selector_loss_alpha=0.0)
         assert 0.0 <= loss <= 1.0
 
+    @pytest.mark.parametrize("loss_type", ["ce", "tv", "lambda"])
+    def test_a_step_with_nothing_to_train_is_gated(self, monkeypatch, loss_type):
+        """Every objective takes the zero loss when the position weights sum to <= 1.
+
+        The one label sits three positions after the anchor, so its decay weight is < 1.
+        """
+        config = _get_dflash2_config(selector_loss_alpha=0.0)
+        config.update(
+            dflash_lk_loss_type=loss_type,
+            dflash_loss_objective="decay",
+            dflash_loss_decay_factor=4.0,
+        )
+        model = get_tiny_llama(num_hidden_layers=4)
+        mtsp.convert(model, [("dflash", config)])
+        model.float().train()
+        monkeypatch.setattr(
+            model,
+            "_sample_anchor_positions",
+            lambda *_: (torch.tensor([[1]]), torch.tensor([[True]])),
+        )
+        input_ids = _make_batch(model.dflash_config.vocab_size)[0][:1]
+        loss_mask = torch.zeros(1, SEQ_LEN)
+        loss_mask[0, 4] = 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
+
     @pytest.mark.parametrize(
         ("loss_type", "distill", "ce_grad"),
         [("ce", False, True), ("tv", False, True), ("lambda", False, True), ("ce", True, False)],
diff --git a/tests/unit/torch/speculative/plugins/test_hf_domino.py b/tests/unit/torch/speculative/plugins/test_hf_domino.py
index 5f0fdf996d3..a159369000f 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_domino.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_domino.py
@@ -31,7 +31,7 @@
 
 import modelopt.torch.speculative as mtsp
 from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG
-from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel
+from modelopt.torch.speculative.plugins.hf_dflash import DraftBatch, HFDFlashModel
 from modelopt.torch.speculative.plugins.hf_domino import HFDominoModel, compute_lambda_base
 from modelopt.torch.speculative.plugins.modeling_dflash import DFlashModule
 from modelopt.torch.speculative.plugins.modeling_domino import DominoModule
@@ -132,10 +132,10 @@ def test_forward_produces_dual_loss_and_grads(self):
         assert out.loss.requires_grad
         assert out.loss.dim() == 0
         # dual-loss bookkeeping
-        assert "base_loss" in out.domino_metrics
-        assert "final_loss" in out.domino_metrics
-        assert "base_accuracy" in out.domino_metrics
-        assert out.domino_metrics["lambda_base"] == 1.0  # default before any callback
+        assert "loss_backbone" in out.draft_metrics
+        assert "loss_final" in out.draft_metrics
+        assert "backbone_accuracy" in out.draft_metrics
+        assert out.draft_metrics["lambda_base"] == 1.0  # default before any callback
 
         out.loss.backward()
         gru_grad = model.dflash_module.prefix_gru.weight_ih_l0.grad
@@ -146,7 +146,7 @@ def test_forward_produces_dual_loss_and_grads(self):
         assert backbone_grad is not None and torch.isfinite(backbone_grad).all()
 
     def test_lambda_zero_uses_final_only(self):
-        """With lambda_base=0 the loss equals final_loss (correction head only)."""
+        """With lambda_base=0 the loss equals loss_final (correction head only)."""
         model = get_tiny_llama(num_hidden_layers=4)
         mtsp.convert(model, [("dflash", _get_domino_config())])
         model.train()
@@ -155,7 +155,7 @@ def test_lambda_zero_uses_final_only(self):
         vocab = model.dflash_config.vocab_size
         input_ids, attention_mask, labels = self._make_batch(vocab)
         out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
-        assert abs(out.loss.item() - out.domino_metrics["final_loss"]) < 1e-4
+        assert abs(out.loss.item() - out.draft_metrics["loss_final"].item()) < 1e-4
 
 
 class TestDominoSwa:
@@ -213,15 +213,16 @@ def test_drafts_are_the_training_head_argmax(self, pure_draft_prefix_len):
 
         with torch.no_grad():
             _, hidden_states = model._generation_base(input_ids)
-            hidden = model._draft_block(base_token, hidden_states)
-            final_logits = model._apply_domino_head(
-                hidden,
-                model._base_model_lm_head(hidden),
+            batch = DraftBatch(
                 torch.cat([base_token, drafts], dim=1),  # the anchor, then the drafts
-                torch.zeros(2, 1, dtype=torch.long),
-                1,
+                loss_mask=None,
+                anchors=torch.zeros(2, 1, dtype=torch.long),
+                keep=None,
+                base_outputs=None,
+                hidden=model._draft_block(base_token, hidden_states),
             )
-        assert torch.equal(final_logits.argmax(dim=-1), drafts)
+            final_logits = model._draft_logits(batch).final  # [B, 1 block, bs, V]
+        assert torch.equal(final_logits[:, 0].argmax(dim=-1), drafts)
 
 
 class TestLambdaSchedule:
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py
index 66fc64d37db..2e2c541a090 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py
@@ -37,7 +37,7 @@
 import modelopt.torch.speculative as mtsp
 from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG
 from modelopt.torch.speculative.plugins.dflash_losses import _tvd_chunk, _tvd_per_token
-from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel
+from modelopt.torch.speculative.plugins.hf_dflash import DraftBatch, HFDFlashModel
 from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel
 from modelopt.torch.speculative.plugins.modeling_dflash import (
     DFlashBaseModelOutput,
@@ -174,9 +174,9 @@ def test_forward_loss_metrics_and_grads(self, head_type):
         assert out.loss.requires_grad
         assert out.loss.dim() == 0
         # three-term loss bookkeeping
-        for key in ("ce_loss", "l1_loss", "confidence_loss", "base_accuracy"):
-            assert key in out.dspark_metrics
-        assert out.dspark_metrics["confidence_loss"] == 0.0  # no confidence head here
+        for key in ("loss_ce", "loss_tvd", "confidence_loss", "backbone_accuracy"):
+            assert key in out.draft_metrics
+        assert out.draft_metrics["confidence_loss"] == 0.0  # no confidence head here
 
         out.loss.backward()
         head_grad = model.dflash_module.markov_w2.weight.grad
@@ -196,7 +196,7 @@ def test_confidence_head_contributes_grads(self):
 
         input_ids, attention_mask, labels = self._make_batch(model.dflash_config.vocab_size)
         out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
-        assert out.dspark_metrics["confidence_loss"] > 0.0
+        assert out.draft_metrics["confidence_loss"] > 0.0
 
         out.loss.backward()
         conf_grad = model.dflash_module.confidence_proj.weight.grad
@@ -1063,25 +1063,15 @@ def test_batch_with_no_valid_anchor(self):
     def test_loss_branch_with_zero_total_weight(self):
         """Anchors exist but every label position is masked, so the three terms are skipped."""
         model = self._model()
-        m = model.dflash_module
         n_blocks, bsz = 2, 1
-        input_ids, anchors, _, args = _draft_args(model, n_blocks=n_blocks, bsz=bsz)
-        hidden = m(*args)
-        vocab = model.dflash_config.vocab_size
-        backbone_logits = torch.randn(bsz, n_blocks * BLOCK_SIZE, vocab, dtype=hidden.dtype)
-        base_outputs = DFlashBaseModelOutput(args[1])  # no base distribution: not read here
-        final_logits, confidence_logits = model._apply_markov_head(
-            hidden, backbone_logits, input_ids, anchors, n_blocks
-        )
-        loss, _, _ = model._compute_dspark_loss(
-            backbone_logits,
-            final_logits,
-            confidence_logits,
+        input_ids, anchors, keep, args = _draft_args(model, n_blocks=n_blocks, bsz=bsz)
+        batch = DraftBatch(
             input_ids,
+            torch.zeros(bsz, SEQ_LEN),  # blocks are kept, but no label position carries weight
             anchors,
-            torch.ones(bsz, n_blocks),  # blocks are kept ...
-            torch.zeros(bsz, SEQ_LEN),  # ... but no label position carries weight
-            base_outputs,
+            keep,
+            DFlashBaseModelOutput(args[1]),  # no base distribution: not read here
+            model.dflash_module(*args),
         )
-        loss.backward()
+        model._compute_loss(model._loss_context(batch)).loss.backward()
         assert not self._ungraded(model), f"no gradient for {self._ungraded(model)}"

From 29a932b0d4ea37ab2d52e28d3b78f4dcb0cae395 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 13:56:57 +0000
Subject: [PATCH 5/6] test(speculative): pin how each shared loss term reduces

Before the terms were shared, DFlash2's _lk_loss carried its own copy of the cross-entropy
reduction, so test_zero_decay_and_unit_scale_is_exactly_cross_entropy compared two copies.
Now the lambda blend calls the ce term, and that test compares the term with itself. The
kd, tvd and tv reductions were never compared with anything. A parametrized test now runs
ce, kd, tvd and tv through the real _loss_context and _compute_loss under decay weights, and
compares each loss bitwise with a verbatim copy of the old reduction:
HFDFlashModel._compute_loss for ce and kd, DSpark's _compute_dspark_loss for tvd, and
DFlash2's _lk_loss for tv.

The terms' needs_ce_grad declarations need no new test. They now set the per-token
cross-entropy's grad mode, which test_cross_entropy_gradient_follows_the_objective already
pins for ce, tv, lambda and KD. If the tv or lambda declaration is flipped, the loss value
does not change, but the backbone trains through the selector alone; the shipped
dflash2.yaml uses lambda with the selector.

Mutation check. Each source change below passes the rest of tests/unit/torch/speculative
and test_hf_spec_rope_export.py (519 passed), and fails exactly the case named:
- ce, kd, tvd or tv divides by the supervised count instead of the weight sum: the new
  test's [ce], [kd], [tvd] or [tv]. Reducing ce under the mask instead of the weights also
  fails [ce].
- tv or lambda declares needs_ce_grad = False, or kd declares it True:
  test_cross_entropy_gradient_follows_the_objective's [tv], [lambda] or [kd].
Flipping ce's own declaration fails four existing tests and that test's [ce].

Each new case takes under 0.1 s on 4 CPUs.

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     | 42 +++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index 57487aaa522..ec86aa115c7 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -26,6 +26,7 @@
 
 import pytest
 import torch
+import torch.nn.functional as F
 from _test_utils.torch.speculative.dflash import get_dflash_config
 from _test_utils.torch.transformers_models import (
     get_tiny_llama,
@@ -41,8 +42,11 @@
 import modelopt.torch.speculative.plugins.hf_dflash as hf_dflash
 from modelopt.torch.speculative.plugins.dflash_losses import (
     DraftLogits,
+    _block_position_weights,
     _block_targets,
     _dpace_position_weights,
+    _tvd_per_token,
+    objective_term,
 )
 from modelopt.torch.speculative.plugins.hf_dflash import (
     DFlashAttention,
@@ -262,6 +266,22 @@ def test_convert_dpace_with_decay_factor_warns(self, caplog):
         assert any("dflash_loss_decay_factor" in r.message for r in caplog.records)
 
 
+def _legacy_term_loss(name, flat_logits, flat_targets, flat_teacher, flat_weights):
+    """Verbatim copy of the divergences' pre-refactor reductions, kept as the semantic reference."""
+    valid_count = flat_weights.sum() + 1e-6
+    if name == "kd":
+        target_soft = torch.softmax(flat_teacher, dim=-1)
+        draft_logsoft = torch.log_softmax(flat_logits, dim=-1)
+        per_token = -(target_soft * draft_logsoft).sum(dim=-1)
+    elif name == "tvd":
+        per_token = _tvd_per_token(flat_logits, flat_teacher)
+    else:
+        per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none")
+        if name == "tv":
+            per_token = 1.0 - torch.exp(-per_token)
+    return (per_token * flat_weights).sum() / valid_count
+
+
 class TestDPaceLossIntegration:
     """Exercise the _compute_loss block-weighting branches on CPU."""
 
@@ -316,6 +336,28 @@ def test_kd_without_a_teacher_raises(self):
         with pytest.raises(ValueError, match="base model's distribution"):
             _compute_loss(model, *self._make_inputs(), base_outputs=DFlashBaseModelOutput(None))
 
+    @pytest.mark.parametrize("name", ["ce", "kd", "tvd", "tv"])
+    def test_each_term_reduces_as_before(self, name):
+        """Under non-uniform position weights, each divergence is bitwise its legacy reduction."""
+        vocab = 32
+        model = self._converted_model("decay", dflash_loss_decay_factor=4.0)
+        model._objective = [objective_term(name, 1.0)]
+        logits, input_ids, anchors, keep, loss_mask = self._make_inputs(vocab=vocab)
+        base_outputs = DFlashBaseModelOutput(None, logits=torch.randn(1, SEQ_LEN, vocab))
+        loss, _ = _compute_loss(model, logits, input_ids, anchors, keep, loss_mask, base_outputs)
+
+        targets = _block_targets(input_ids, anchors, keep, loss_mask, BLOCK_SIZE, 0)
+        weights = _block_position_weights(targets.mask, "decay", 1, 4.0, model.dflash_dpace_alpha)
+        teacher = model._teacher_logits(base_outputs, targets.positions)
+        legacy = _legacy_term_loss(
+            name,
+            logits.view(-1, vocab),
+            targets.token_ids.view(-1),
+            teacher.reshape(-1, vocab),
+            weights.view(-1),
+        )
+        assert torch.equal(loss, legacy)
+
 
 class TestDFlashSaveRestore:
     """Test DFlash model save and restore."""

From 28f4f45d34f9d8df93ae0b16a2c74cb97768665a Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Tue, 13 Oct 2026 14:18:42 +0000
Subject: [PATCH 6/6] docs(speculative): extend the draft_metrics entry to
 DSpark and Domino

DSpark's dspark_metrics and Domino's domino_metrics have been in every release since 0.46.0.
The DSpark and Domino commits remove both keys, rename three inner keys of each, and return
tensors where both returned Python floats. Nothing in the repo reads these keys, but user
code that reads them breaks.

The backward-breaking entry for DFlash2's and LiLiCorr's keys now also names these two keys
and the renamed inner keys.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 CHANGELOG.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 294bcc88e82..2907f626226 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -36,7 +36,7 @@ Changelog
 
 - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``.
 - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata.
-- The DFlash2 and LiLiCorr training forwards now return their metrics in one ``draft_metrics`` dict of detached 0-dim tensors (call ``.item()`` for a Python float), replacing ``selector_metrics`` and ``lilicorr_metrics``.
+- The DFlash2, LiLiCorr, DSpark and Domino training forwards now return their metrics in one ``draft_metrics`` dict of detached 0-dim tensors (call ``.item()`` for a Python float), replacing ``selector_metrics``, ``lilicorr_metrics``, ``dspark_metrics`` and ``domino_metrics``. DSpark's ``ce_loss``/``l1_loss``/``base_accuracy`` are now ``loss_ce``/``loss_tvd``/``backbone_accuracy``, and Domino's ``final_loss``/``base_loss``/``base_accuracy`` are now ``loss_final``/``loss_backbone``/``backbone_accuracy``.
 
 **Deprecations**