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**
 
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_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):
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index 0891f9a924b..69b5fdc3f1f 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -64,18 +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_chunk,
-    _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
@@ -100,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(
@@ -108,22 +99,12 @@ 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(
                 "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."""
@@ -131,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.
-
-        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
+    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),
+        ]
 
-        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.
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)
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."""
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)}"