From 73a7bb851bca2e7243cfb840701b1df897ccebcb Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 14:53:29 +0000
Subject: [PATCH 1/4] refactor(speculative): one online base forward for the
 DFlash family

The online teacher depended on the variant: DFlash, DFlash2 and LiLiCorr ran the CausalLM
forward and gathered their teacher rows out of its full [B, seq, vocab] logits, DSpark ran the
inner model and projected only the rows its loss reads, and Domino kept no teacher at all.
Offline and streaming training already used the row projection for everyone.

_online_base_outputs is no longer a variant hook, and every batch gets one teacher:
_teacher_logits projects the base's final hidden state, the inner model's output, at just the
rows a loss reads, as it does offline.
- A text batch runs the inner model once and keeps its final hidden state next to the target
  hidden states.
- A multimodal batch keeps the top-level forward, which inserts the image/video features, and
  a forward hook on the inner model reads its output during it. The hook exists for that
  forward only and is removed in a finally. It reads where a text batch reads, not at lm_head's
  input, which a CausalLM may rescale first (Inkling divides it by a multiplier), so the two
  batches cannot get different teachers. The inner model must run once over the whole
  sequence, or the step raises; logits_to_keep narrows only lm_head's input.
- DSpark's and Domino's overrides go, so their image batches reach the multimodal forward too.
  Theirs ran the language model on the raw token ids and dropped the image features.
- The multimodal forward now also gets image_position_ids and video_position_ids, without which
  Gemma 4's vision tower fails: DFlash's Gemma 4 image and video batches failed there before,
  and DSpark's and Domino's would have started to. It still gets no token_type_ids, from which
  Gemma 3 and PaliGemma build their bidirectional attention over the image, so on transformers
  5.5 their image batches raise, now for DSpark and Domino too, and on 5.12 and later they
  attend causally over the image, as DFlash's did.

Behavior change, for online DFlash, DFlash2 and LiLiCorr only: their teacher is lm_head on the
final hidden state instead of the CausalLM's logits. That equals the old teacher up to matmul
rounding, except on a base whose CausalLM rescales its logits around lm_head, which the teacher
now leaves out: Gemma's final_logit_softcapping, for which the next commit adds a switch, and
the scaling of Granite (logits_scaling), Cohere (logit_scale), Falcon-H1 (lm_head_multiplier)
and Inkling (logits_mup_width_multiplier), which changes the teacher's temperature by that
factor. Text batches no longer build full-sequence logits, an image batch's are dropped, and
the training output's logits are None.

Tested: tests/unit/torch/speculative and tests/unit/torch/export/test_hf_spec_rope_export.py
533 passed on transformers 5.5.4. The online case of the full-sequence test fails on
the base tree, where the CausalLM forward fed lm_head the whole sequence; a new test pins the
projected online teacher to the CausalLM's logits. On a tiny Qwen3-VL, new tests check that an
image batch's teacher equals the model's own logits at the rows read, that an image batch and
a text batch read the same final hidden state before the image even when the CausalLM
rescales lm_head's input, that another image changes DFlash's, DSpark's and Domino's base rows
from the image on (on the base tree it left DSpark's and Domino's unchanged), that the hook
never exists for a text batch and is gone after an image batch, also one whose forward fails,
that logits_to_keep leaves the teacher whole, and that an inner model run twice or on fewer
rows raises. On a tiny Gemma 4 VLM, a new test checks that an image and a video reach the base.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../torch/speculative/plugins/hf_dflash.py    |  55 ++--
 .../torch/speculative/plugins/hf_domino.py    |  11 -
 .../torch/speculative/plugins/hf_dspark.py    |  15 -
 .../speculative/plugins/modeling_dflash.py    |   7 +-
 .../speculative/plugins/test_hf_dflash.py     | 265 +++++++++++++++++-
 5 files changed, 294 insertions(+), 59 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index a0c2cc34bb1..02a895bf935 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -120,6 +120,8 @@
         "pixel_values_videos",
         "image_grid_thw",
         "video_grid_thw",
+        "image_position_ids",
+        "video_position_ids",
         "mm_token_type_ids",
         "image_sizes",
         "images",
@@ -853,13 +855,32 @@ def _online_base_outputs(
         cache_position=None,
         **kwargs,
     ):
-        """Run the frozen base model: the target hidden states, and its logits for KD."""
-        # Multimodal models need the top-level conditional-generation forward so their
-        # image/video features are inserted before the language model runs.  Keep the
-        # long-standing narrow call for text-only models.
+        """Run the frozen base model: the target hidden states, and its final hidden state.
+
+        The final hidden state is the inner model's output for every batch, and
+        ``_teacher_logits`` projects it only at the rows the loss reads. A text batch runs the
+        inner model once. A multimodal batch needs the top-level forward to insert its
+        image/video features, and a hook on the inner model reads its output during it.
+        """
         base_forward_kwargs = _multimodal_forward_kwargs(kwargs)
-        with torch.no_grad():
-            if base_forward_kwargs:
+        if not base_forward_kwargs:
+            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),
+                base_hidden=base_out.last_hidden_state,
+            )
+
+        # Read where a text batch reads it, not at lm_head's input, which a CausalLM may rescale
+        # first (Inkling divides it by a multiplier).
+        base_hidden = []
+        hook = self._base_model.register_forward_hook(
+            lambda _module, _args, output: base_hidden.append(output[0])
+        )
+        try:
+            with torch.no_grad():
                 raw_outputs = super().forward(
                     input_ids=input_ids,
                     attention_mask=attention_mask,
@@ -873,20 +894,22 @@ def _online_base_outputs(
                     return_dict=True,
                     **base_forward_kwargs,
                 )
-            else:
-                raw_outputs = super().forward(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    output_hidden_states=True,
-                )
-
+        finally:
+            hook.remove()
         if not getattr(raw_outputs, "hidden_states", None):
             raise RuntimeError(
                 "The base model did not return hidden states required for DFlash training. "
                 "Ensure its top-level multimodal forward supports output_hidden_states=True."
             )
+        if len(base_hidden) != 1 or base_hidden[0].shape[:2] != input_ids.shape:
+            raise RuntimeError(
+                "The teacher reads the base model's final hidden state off its inner model, so the "
+                "multimodal forward must run that once over the whole sequence; it ran it on "
+                f"{[tuple(h.shape) for h in base_hidden]} for input_ids {tuple(input_ids.shape)}."
+            )
         return DFlashBaseModelOutput(
-            target_hidden=self._target_hidden(raw_outputs.hidden_states), logits=raw_outputs.logits
+            target_hidden=self._target_hidden(raw_outputs.hidden_states),
+            base_hidden=base_hidden[0],
         )
 
     def _draft_logits(self, batch):
@@ -974,8 +997,8 @@ def forward(
         - Label alignment: position k predicts the token at anchor + k + ``_label_shift``
         - Position weighting by D-PACE or exponential decay
 
-        A variant changes the step only through its hooks: ``_online_base_outputs``,
-        ``_draft_logits`` (its head), ``_position_weighting`` and ``_compute_loss``.
+        A variant changes the step only through its hooks: ``_draft_logits`` (its head),
+        ``_position_weighting`` and ``_compute_loss``.
         """
         if not self.training:
             if self.dflash_offline:
diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py
index 0a4bc9fdbd4..2dd4cecfd94 100644
--- a/modelopt/torch/speculative/plugins/hf_domino.py
+++ b/modelopt/torch/speculative/plugins/hf_domino.py
@@ -67,7 +67,6 @@
 from ..dflash.conversion import DominoDMRegistry
 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
 
 logger = logging.getLogger(__name__)
@@ -143,16 +142,6 @@ def _position_weighting(self):
         # dflash_loss_objective is not wired for Domino yet: it always decays.
         return "decay"
 
-    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
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index 69b5fdc3f1f..a4dfbe3032e 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -68,7 +68,6 @@
 from ..dflash.conversion import DSparkDMRegistry
 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
 
 logger = logging.getLogger(__name__)
@@ -124,20 +123,6 @@ def _position_weighting(self):
         # dflash_loss_objective is not wired for DSpark yet: it always decays.
         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():
-            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),
-            base_hidden=base_out.last_hidden_state,
-        )
-
     def _head_inputs(self, batch):
         """Each block position's teacher-forced previous token, and the draft's hidden states.
 
diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py
index 45870d8b6c7..836234fc052 100644
--- a/modelopt/torch/speculative/plugins/modeling_dflash.py
+++ b/modelopt/torch/speculative/plugins/modeling_dflash.py
@@ -118,10 +118,9 @@ def _get_sink_attention_fn():
 class DFlashBaseModelOutput:
     """What a DFlash-family draft takes from the base model in a training step.
 
-    ``logits`` holds the full base logits when the producer already has them (DFlash's
-    online CausalLM forward, or an offline batch carrying ``base_model_logits``);
-    otherwise ``base_hidden`` stays unprojected, and ``HFDFlashModel._teacher_logits``
-    projects only the rows a loss reads.
+    ``logits`` holds the full base logits when an offline batch carries them
+    (``base_model_logits``); otherwise ``base_hidden`` stays unprojected, and
+    ``HFDFlashModel._teacher_logits`` projects only the rows a loss reads.
     """
 
     target_hidden: torch.Tensor  # concatenated hidden states from target layers [B, seq, N*H]
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index ec86aa115c7..6dd54cef812 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -30,10 +30,11 @@
 from _test_utils.torch.speculative.dflash import get_dflash_config
 from _test_utils.torch.transformers_models import (
     get_tiny_llama,
+    get_tiny_qwen3vl,
     tf_modelopt_state_and_output_tester,
 )
 from torch.nn.attention.flex_attention import create_mask
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, Gemma4Config, Gemma4ForConditionalGeneration
 from transformers.modeling_layers import GradientCheckpointingLayer
 
 import modelopt.torch.opt as mto
@@ -1714,11 +1715,12 @@ def test_kd_loss_matches_full_logits(self):
         )
         torch.testing.assert_close(loss, loss_full)
 
-    def test_offline_step_never_projects_the_full_sequence(self):
-        """Offline/streaming with KD on, lm_head only ever sees draft rows and teacher rows."""
+    @pytest.mark.parametrize("offline", [False, True], ids=["online", "offline"])
+    def test_step_never_projects_the_full_sequence(self, offline):
+        """With KD on, lm_head only ever sees draft rows and teacher rows."""
         model = get_tiny_llama(num_hidden_layers=4)
         model.config.num_orig_hidden_layers = 4
-        mtsp.convert(model, [("dflash", get_dflash_config(offline=True))])
+        mtsp.convert(model, [("dflash", get_dflash_config(offline=offline))])
         assert model.dflash_self_logit_distillation
         model.train()
         seen = []
@@ -1727,17 +1729,254 @@ def test_offline_step_never_projects_the_full_sequence(self):
         )
         bsz, hidden = 2, model.config.hidden_size
         dtype = next(model.dflash_module.parameters()).dtype
-        base_model_outputs = {
-            "aux_hidden_states": torch.randn(
-                bsz, SEQ_LEN, len(model.target_layer_ids) * hidden, dtype=dtype
-            ),
-            "base_model_hidden_states": torch.randn(bsz, SEQ_LEN, hidden, dtype=dtype),
-        }
+        kwargs = {}
+        if offline:
+            kwargs["base_model_outputs"] = {
+                "aux_hidden_states": torch.randn(
+                    bsz, SEQ_LEN, len(model.target_layer_ids) * hidden, dtype=dtype
+                ),
+                "base_model_hidden_states": torch.randn(bsz, SEQ_LEN, hidden, dtype=dtype),
+            }
         input_ids = torch.randint(1, model.config.vocab_size, (bsz, SEQ_LEN))
         model(
-            input_ids=input_ids,
-            attention_mask=torch.ones_like(input_ids),
-            base_model_outputs=base_model_outputs,
+            input_ids=input_ids, attention_mask=torch.ones_like(input_ids), **kwargs
         ).loss.backward()
         assert len(seen) == 2, seen  # the draft's logits, then the KD teacher rows
         assert (bsz, SEQ_LEN, hidden) not in seen, seen
+
+    def test_online_teacher_is_the_causal_lm_distribution(self):
+        """The rows projected from the inner model's final hidden are the CausalLM's logits."""
+        model = self._model().float()
+        labels = self._labels()
+        input_ids = torch.randint(1, model.config.vocab_size, (2, SEQ_LEN))
+        attention_mask = torch.ones_like(input_ids)
+        outputs = model._online_base_outputs(input_ids, attention_mask)
+        model.eval()
+        with torch.no_grad():
+            full = model(input_ids=input_ids, attention_mask=attention_mask).logits
+        torch.testing.assert_close(
+            model._teacher_logits(outputs, labels), self._rows(full, labels - 1)
+        )
+
+
+# Qwen3-VL's vision placeholder ids, moved inside get_tiny_qwen3vl's 32-token vocab.
+_VISION_START, _IMAGE = 28, 30
+
+
+def _tiny_vlm(arch=None):
+    """A tiny Qwen3-VL converted online, in fp32; ``arch`` updates the draft's architecture."""
+    model = get_tiny_qwen3vl(num_hidden_layers=4)
+    model.config.vision_start_token_id, model.config.image_token_id = _VISION_START, _IMAGE
+    config = get_dflash_config()
+    config["dflash_architecture_config"].update(arch or {})
+    mtsp.convert(model, [("dflash", config)])
+    return model.float()
+
+
+def _image_batch():
+    """Two sequences, each holding a 2x2-patch image at positions 4-7, and its pixels."""
+    g = torch.Generator().manual_seed(0)
+    input_ids = torch.randint(1, _VISION_START, (2, SEQ_LEN), generator=g)
+    input_ids[:, 3] = _VISION_START
+    input_ids[:, 4:8] = _IMAGE
+    return {
+        "input_ids": input_ids,
+        "attention_mask": torch.ones_like(input_ids),
+        # One row per 4x4 patch of 3 channels: get_tiny_qwen3vl merges no patches.
+        "pixel_values": torch.randn(2 * 4, 3 * 4 * 4, generator=g),
+        "image_grid_thw": torch.tensor([[1, 2, 2], [1, 2, 2]]),
+        "mm_token_type_ids": (input_ids == _IMAGE).long(),
+    }
+
+
+def _tiny_gemma4_and_batch(modality):
+    """A tiny Gemma 4 VLM converted online, in fp32, and a batch with an image or a video at 3-6."""
+    torch.manual_seed(0)
+    image_token, video_token = 62, 63
+    text_config = {
+        "vocab_size": 64,
+        "hidden_size": 32,
+        "intermediate_size": 32,
+        "num_hidden_layers": 4,
+        "num_attention_heads": 4,
+        "num_key_value_heads": 2,
+        "head_dim": 8,
+        "max_position_embeddings": 64,
+        "sliding_window": 8,
+        "vocab_size_per_layer_input": 64,
+        "hidden_size_per_layer_input": 8,
+    }
+    vision_config = {
+        "hidden_size": 16,
+        "intermediate_size": 16,
+        "num_hidden_layers": 1,
+        "num_attention_heads": 2,
+        "num_key_value_heads": 2,
+        "head_dim": 8,
+        "patch_size": 4,
+        "position_embedding_size": 64,
+    }
+    model = Gemma4ForConditionalGeneration(
+        Gemma4Config(
+            text_config=text_config,
+            vision_config=vision_config,
+            image_token_id=image_token,
+            # The other special tokens, inside the tiny vocabulary too.
+            boi_token_id=60,
+            eoi_token_id=61,
+            video_token_id=video_token,
+            audio_token_id=57,
+            boa_token_id=59,
+            eoa_token_index=58,
+        )
+    )
+    mtsp.convert(model, [("dflash", get_dflash_config())])
+    g = torch.Generator().manual_seed(1)
+    input_ids = torch.randint(1, 56, (2, SEQ_LEN), generator=g)
+    # 6x6 patches of 4x4 pixels, pooled to 2x2 tokens; a video has one such frame.
+    input_ids[:, 3:7] = image_token if modality == "image" else video_token
+    ys, xs = torch.meshgrid(torch.arange(6), torch.arange(6), indexing="ij")
+    pixels = torch.randn(2, 6 * 6, 3 * 4 * 4, generator=g)
+    positions = torch.stack([xs.flatten(), ys.flatten()], -1).expand(2, -1, -1)
+    if modality == "image":
+        visual = {"pixel_values": pixels, "image_position_ids": positions}
+    else:
+        visual = {"pixel_values_videos": pixels[:, None], "video_position_ids": positions[:, None]}
+    batch = {
+        "input_ids": input_ids,
+        "attention_mask": torch.ones_like(input_ids),
+        **visual,
+        # Gemma 4's processor marks image tokens 1 and video tokens 2.
+        "mm_token_type_ids": (input_ids == image_token).long() + 2 * (input_ids == video_token),
+    }
+    return model.float(), batch
+
+
+class TestMultimodalTeacher:
+    """An image batch gets a text batch's teacher: lm_head projecting the inner model's output.
+
+    The top-level forward inserts the image, and a hook reads the inner model's output during it.
+    """
+
+    @staticmethod
+    def _labels():
+        torch.manual_seed(0)
+        return torch.randint(1, SEQ_LEN, (2, 3, BLOCK_SIZE))
+
+    @staticmethod
+    def _own_rows(model, batch, positions):
+        """Rows of the model's own logits, which for Qwen3-VL are lm_head's output as it is."""
+        model.eval()
+        with torch.no_grad():
+            logits = model(**batch).logits
+        return logits[torch.arange(positions.shape[0]).view(-1, 1, 1), positions]
+
+    def test_teacher_projects_the_final_hidden_state_over_the_whole_sequence(self):
+        model = _tiny_vlm()
+        batch = _image_batch()
+        lm_head_inputs = []
+        model._base_model_lm_head.register_forward_pre_hook(
+            lambda _module, args: lm_head_inputs.append(args[0])
+        )
+        outputs = model._online_base_outputs(**batch)
+        assert outputs.logits is None
+        # Qwen3-VL feeds the inner model's output to lm_head as it is.
+        assert torch.equal(outputs.base_hidden, lm_head_inputs[0])
+        labels = self._labels()
+        torch.testing.assert_close(
+            model._teacher_logits(outputs, labels), self._own_rows(model, batch, labels - 1)
+        )
+
+    def test_text_and_image_batches_read_the_same_final_hidden_state(self):
+        """Up to the image, also when the CausalLM rescales lm_head's input, as Inkling does."""
+        model = _tiny_vlm()
+        model._base_model_lm_head.register_forward_pre_hook(lambda _module, args: (args[0] / 24,))
+        batch = _image_batch()
+        image = model._online_base_outputs(**batch).base_hidden
+        text = model._online_base_outputs(batch["input_ids"], batch["attention_mask"]).base_hidden
+        torch.testing.assert_close(image[:, :4], text[:, :4])
+
+    @pytest.mark.parametrize(
+        "arch",
+        [
+            {},
+            {"projector_type": "dspark", "markov_rank": 16},
+            {"projector_type": "domino", "emb_dim": 16, "gru_hidden_dim": 32},
+        ],
+        ids=["dflash", "dspark", "domino"],
+    )
+    def test_every_variant_trains_on_the_image(self, arch):
+        """Another image in the same text changes the base's rows from the image on, not before."""
+        model = _tiny_vlm(arch)
+        batch = _image_batch()
+        other = {**batch, "pixel_values": torch.randn_like(batch["pixel_values"])}
+        first, second = (model._online_base_outputs(**b) for b in (batch, other))
+        for name in ("target_hidden", "base_hidden"):
+            a, b = getattr(first, name), getattr(second, name)
+            torch.testing.assert_close(a[:, :4], b[:, :4], msg=name)
+            assert not torch.allclose(a[:, 4:], b[:, 4:]), name
+        model.train()
+        out = model(**batch)
+        assert torch.isfinite(out.loss)
+        out.loss.backward()
+
+    @pytest.mark.parametrize("modality", ["image", "video"])
+    def test_a_gemma4_image_or_video_reaches_the_base(self, modality):
+        """Gemma 4's vision tower reads the patch positions its processor sends with the pixels."""
+        model, batch = _tiny_gemma4_and_batch(modality)
+        pixels = "pixel_values" if modality == "image" else "pixel_values_videos"
+        other = {**batch, pixels: torch.randn_like(batch[pixels])}
+        first, second = (model._online_base_outputs(**b).base_hidden for b in (batch, other))
+        torch.testing.assert_close(first[:, :3], second[:, :3])
+        assert not torch.allclose(first[:, 3:], second[:, 3:])
+
+    def test_the_inner_model_is_hooked_only_inside_an_image_batch_forward(self):
+        """Never for a text batch; an image batch's hook goes even when its forward fails."""
+        model = _tiny_vlm()
+        base = model._base_model
+        hooks_seen = []
+        base.register_forward_pre_hook(
+            lambda module, _args: hooks_seen.append(len(module._forward_hooks))
+        )
+        batch = _image_batch()
+        model._online_base_outputs(batch["input_ids"], batch["attention_mask"])
+        model._online_base_outputs(**batch)
+        assert hooks_seen == [0, 1]
+        assert not base._forward_hooks
+        with pytest.raises(RuntimeError):  # pixels that do not fit the vision tower
+            model._online_base_outputs(**{**batch, "pixel_values": torch.randn(3, 5)})
+        assert not base._forward_hooks
+
+    def test_logits_to_keep_does_not_narrow_the_teacher(self, monkeypatch):
+        """It narrows lm_head's input, not the inner model's output that the teacher reads."""
+        model = _tiny_vlm()
+        batch = _image_batch()
+        want = model._online_base_outputs(**batch).base_hidden
+        monkeypatch.setattr(
+            hf_dflash,
+            "_MULTIMODAL_FORWARD_KWARGS",
+            hf_dflash._MULTIMODAL_FORWARD_KWARGS | {"logits_to_keep"},
+        )
+        assert torch.equal(model._online_base_outputs(**batch, logits_to_keep=1).base_hidden, want)
+
+    @pytest.mark.parametrize("run", ["twice", "short"])
+    def test_the_inner_model_must_run_once_over_the_whole_sequence(self, run):
+        """Otherwise the forward gives the teacher no one final hidden state for every row."""
+        model = _tiny_vlm()
+        batch = _image_batch()
+        base = model._base_model
+
+        def run_it_again(_module, _args):
+            base(input_ids=batch["input_ids"])
+
+        def drop_the_last_row(_module, _args, output):
+            output.last_hidden_state = output.last_hidden_state[:, :-1]
+
+        if run == "twice":
+            model._base_model_lm_head.register_forward_pre_hook(run_it_again)
+        else:
+            base.register_forward_hook(drop_the_last_row)
+        hooks = dict(base._forward_hooks)
+        with pytest.raises(RuntimeError, match="once over the whole sequence"):
+            model._online_base_outputs(**batch)
+        assert base._forward_hooks == hooks

From 3bdb736f5693db6d080725a836142c825fb5d62d Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 14:58:07 +0000
Subject: [PATCH 2/4] feat(speculative): opt-in final_logit_softcapping for the
 DFlash-family teacher

The teacher is now lm_head on the base's final hidden state for every batch, which leaves out
what a base's CausalLM does to its logits around lm_head. Gemma-family models cap them, and some
bases need that cap in the teacher, so it comes back as a switch, off by default so that no
other run changes. The logit scaling of other bases stays out, and conversion says so.

- dflash_teacher_softcap: true makes _teacher_logits apply cap * tanh(logits / cap), with cap
  the base config's final_logit_softcapping, on the full rows and on the token_ids path
  LiLiCorr reads, for text and image batches alike. It runs in the base model's own op order.
  Logits handed over in an offline batch are the model's own and are still only gathered.
  Converting raises when the switch is on and the base config sets no
  final_logit_softcapping.
- The switch alone decides the cap. A multimodal base gets the same rule as a text-only one,
  and restore reads the switch from the saved config like any field, so a checkpoint saved
  before the switch existed restores with it off.
- One table, next to the base-model path constants in modeling_fakebase.py, lists the base
  config fields by which a CausalLM scales its logits at lm_head: logits_scaling (Granite
  family, HyperCLOVA X, MiniCPM3), logit_scale (Cohere family), lm_head_multiplier
  (Falcon-H1), output_multiplier (MuseGlimmer) and logits_mup_width_multiplier (Inkling). When
  a loss term reads the teacher and the base config sets one of them to other than 1.0,
  conversion warns once, on rank 0, that the teacher is lm_head's raw projection and does not
  apply it. Nothing else changes.
- Both are decided at convert time, so whether a loss reads the teacher is decided there too:
  _reads_teacher(config) is stored as _needs_base_logits, and LiLiCorr declares its penalty
  and calibration terms from the config instead of from attributes it sets later.
- FakeBaseConfig carries final_logit_softcapping and whichever of those scaling fields the base
  config sets. Offline training on a fake base, which every offline VLM uses, can then use the
  switch and gets the warning.

Tested: tests/unit/torch/speculative and tests/unit/torch/export/test_hf_spec_rope_export.py
558 passed on transformers 5.5.4. New tests pin the capped teacher on both projection
paths and on a tiny Qwen3-VL image batch, that handed-over logits are not capped again, the
error for a switch on a base without a cap, that a capping base converts with the cap the
switch sets whether multimodal or not, that restore takes the switch from the saved config and
leaves it off for a config saved without it, the fake base carrying the cap and each scaling
field, and the warning: once for a tiny Granite with logits_scaling=8, also offline on a fake
base loaded from it, and not for Llama, a scaling of 1.0, a CE-only objective, or a rank other
than 0. A drafter trained online on a tiny capping Gemma-4 VLM, by the base branch or by
0.47.0, loads through load_vlm_or_llm with the cap off: a resumed DSpark step matches the step
of the tree that trained it bitwise, as does a DFlash CE step (base branch), while a DFlash KD
step distills from the uncapped teacher.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 modelopt/torch/speculative/config.py          |  14 ++
 .../torch/speculative/plugins/hf_dflash.py    |  48 ++++-
 .../torch/speculative/plugins/hf_lilicorr.py  |   9 +-
 .../speculative/plugins/modeling_fakebase.py  |  19 ++
 .../speculative/plugins/test_fakebase.py      |  15 ++
 .../speculative/plugins/test_hf_dflash.py     | 178 +++++++++++++++++-
 6 files changed, 262 insertions(+), 21 deletions(-)

diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py
index 1a53b21e4d4..902f71ef19f 100644
--- a/modelopt/torch/speculative/config.py
+++ b/modelopt/torch/speculative/config.py
@@ -102,6 +102,20 @@ class DFlashConfig(ModeloptBaseConfig):
         default=True, description="Whether to use logit distillation from base model."
     )
 
+    dflash_teacher_softcap: bool = ModeloptField(
+        default=False,
+        description=(
+            "Cap the teacher's logits the way the base model caps its own, "
+            "``cap * tanh(logits / cap)`` with ``cap`` the base config's "
+            "``final_logit_softcapping`` (Gemma family). The teacher is the base's ``lm_head`` "
+            "applied to its final hidden state for every batch, text or multimodal, online or "
+            "offline, which leaves that cap out unless this is set. Other scaling a base applies "
+            "at its ``lm_head``, such as ``logits_scaling`` or ``logit_scale``, stays out; "
+            "conversion warns about it. Requires a base config that sets "
+            "``final_logit_softcapping``. The draft's own logits are never capped."
+        ),
+    )
+
     dflash_loss_decay_factor: float = ModeloptField(
         default=0.0,
         description="Gamma for exponential loss decay weighting (paper Eq.4). "
diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 02a895bf935..02333a77917 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -82,6 +82,8 @@
 from transformers.trainer_pt_utils import LabelSmoother
 from transformers.utils import ModelOutput
 
+from modelopt.torch.utils import warn_rank_0
+
 from ..dflash.conversion import DFlashDMRegistry
 from ..dflash.dflash_model import DFlashModel
 from .dflash_flex_attention import build_draft_block_mask
@@ -105,6 +107,7 @@
     _BASE_MODEL_PATHS,
     _EMBED_TOKENS_PATHS,
     _LM_HEAD_PATHS,
+    _LOGIT_SCALING_FIELDS,
     find_final_norm_path,
 )
 from .modeling_final_norm import _apply_base_final_norm
@@ -228,9 +231,12 @@ def _base_model_norm(self):
         path = getattr(self, "base_model_norm_path", None)
         return self.get_submodule(path) if path else None
 
-    @property
-    def _needs_base_logits(self) -> bool:
-        """Whether the loss reads the target's logits; offline training then reconstructs them."""
+    def _reads_teacher(self, config) -> bool:
+        """Whether a loss term reads the base distribution; offline batches must then carry it.
+
+        Becomes ``_needs_base_logits`` at convert time. A variant that adds a term reading the
+        teacher declares it here, from ``config``.
+        """
         return any(term.needs_teacher for _, _, term in self._objective)
 
     @property
@@ -282,6 +288,14 @@ def modify(self, config):
         self.dflash_fp32_master_weights = config.dflash_fp32_master_weights
 
         base_config = self._base_llm_config
+        base_softcap = getattr(base_config, "final_logit_softcapping", None)
+        if config.dflash_teacher_softcap and base_softcap is None:
+            raise ValueError(
+                "dflash_teacher_softcap=True, but the base config sets no "
+                "final_logit_softcapping to cap the teacher's logits with."
+            )
+        self._teacher_softcap = base_softcap if config.dflash_teacher_softcap else None
+
         # Use Qwen3Config (not generic PretrainedConfig) so rope_parameters is
         # auto-populated from rope_theta. DFlash draft uses Qwen3 components.
         self.dflash_config = _Qwen3Config(**config.dflash_architecture_config)
@@ -462,6 +476,18 @@ def modify(self, config):
             self.dflash_module.compile_body()
 
         self._objective = self._build_objective(config)
+        self._needs_base_logits = self._reads_teacher(config)
+        scaling = {
+            name: getattr(base_config, name)
+            for name in _LOGIT_SCALING_FIELDS
+            if getattr(base_config, name, None) not in (None, 1.0)
+        }
+        if self._needs_base_logits and scaling:
+            warn_rank_0(
+                f"The base config sets {', '.join(f'{k}={v}' for k, v in scaling.items())}, "
+                "which the distillation teacher does not apply: the teacher is lm_head's raw "
+                "projection of the base model's final hidden state."
+            )
         # Compiling fuses the TVD chunk's six vocab-wide elementwise passes.
         self._tvd_chunk_fn = (
             torch.compile(_tvd_chunk, dynamic=False, fullgraph=True)
@@ -819,7 +845,8 @@ def _teacher_logits(self, base_outputs, label_positions, token_ids=None):
         Next-token logits: the ones that predict the token at position ``i`` sit at ``i - 1``.
         Only those rows of the base hidden go through the final norm and lm_head, so no
         full-sequence logits are built from it; ``token_ids`` ([B, ..., k]) narrows the
-        projection to just those vocab entries.
+        projection to just those vocab entries. ``dflash_teacher_softcap`` caps what is
+        projected; logits handed over in ``base_outputs`` are the model's own and only gathered.
         """
         positions = (label_positions - 1).clamp(min=0)
         if base_outputs.logits is not None:
@@ -839,10 +866,15 @@ def _teacher_logits(self, base_outputs, label_positions, token_ids=None):
         rows = _gather_rows(base_outputs.base_hidden, positions).to(lm_head.weight.dtype)
         rows = _apply_base_final_norm(rows, base_outputs.base_hidden_prenorm, self._base_model_norm)
         if token_ids is None:
-            return lm_head(rows)
-        logits = torch.einsum("...h,...kh->...k", rows, lm_head.weight[token_ids])
-        bias = getattr(lm_head, "bias", None)
-        return logits if bias is None else logits + bias[token_ids]
+            logits = lm_head(rows)
+        else:
+            logits = torch.einsum("...h,...kh->...k", rows, lm_head.weight[token_ids])
+            bias = getattr(lm_head, "bias", None)
+            if bias is not None:
+                logits = logits + bias[token_ids]
+        cap = self._teacher_softcap
+        # In the base model's op order, so the cap rounds as the base's own does.
+        return logits if cap is None else torch.tanh(logits / cap) * cap
 
     def _online_base_outputs(
         self,
diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py
index be3100e0278..c4216a66236 100644
--- a/modelopt/torch/speculative/plugins/hf_lilicorr.py
+++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py
@@ -175,13 +175,12 @@ def modify(self, config):
             self.dflash_lilicorr_margin,
         )
 
-    @property
-    def _needs_base_logits(self) -> bool:
+    def _reads_teacher(self, config) -> bool:
         """The distractor penalty and the calibration term both read the target's logits."""
         return (
-            super()._needs_base_logits
-            or self.dflash_lilicorr_w_pen > 0.0
-            or self.dflash_lilicorr_w_cal > 0.0
+            super()._reads_teacher(config)
+            or config.dflash_lilicorr_w_pen > 0.0
+            or config.dflash_lilicorr_w_cal > 0.0
         )
 
     def get_exporter(self):
diff --git a/modelopt/torch/speculative/plugins/modeling_fakebase.py b/modelopt/torch/speculative/plugins/modeling_fakebase.py
index 85a3fe3a9f4..8354ce45eff 100644
--- a/modelopt/torch/speculative/plugins/modeling_fakebase.py
+++ b/modelopt/torch/speculative/plugins/modeling_fakebase.py
@@ -92,6 +92,15 @@ def find_final_norm_path(model) -> str | None:
     "language_model.backbone",
 ]
 _VLM_CONFIG_ATTRS = ["text_config", "llm_config"]
+# Base config fields by which a CausalLM scales its logits at lm_head (1.0: no scaling). The
+# DFlash teacher is lm_head's raw projection and leaves them out, so conversion warns about them.
+_LOGIT_SCALING_FIELDS = (
+    "logits_scaling",  # Granite family, HyperCLOVA X, MiniCPM3
+    "logit_scale",  # Cohere family
+    "lm_head_multiplier",  # Falcon-H1
+    "output_multiplier",  # MuseGlimmer
+    "logits_mup_width_multiplier",  # Inkling
+)
 _SAFETENSORS_INDEX_FILENAME = "model.safetensors.index.json"
 # Single-file safetensors names to try, in order.  Mistral native checkpoints
 # use ``consolidated.safetensors`` instead of the HF-standard ``model.safetensors``.
@@ -118,11 +127,14 @@ def __init__(
         rms_norm_eps=1e-6,
         rope_theta=None,
         final_norm_type=None,
+        final_logit_softcapping=None,
         **kwargs,
     ):
         """Initialize FakeBaseConfig with minimal model configuration parameters."""
         super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
         self.rms_norm_eps = rms_norm_eps
+        # The base's logit cap (Gemma family), for dflash_teacher_softcap.
+        self.final_logit_softcapping = final_logit_softcapping
         # Which self-implemented final-norm class FakeBaseModel builds, or None to build no norm
         # (model whose final-norm type we don't know). See _FINAL_NORM_CLASSES /
         # _FINAL_NORM_TYPE_BY_MODEL_TYPE. Persisted so a reloaded config rebuilds the same norm.
@@ -242,6 +254,13 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM
             final_norm_type=_select_final_norm_type(
                 getattr(base_cfg, "model_type", None), base_cfg
             ),
+            final_logit_softcapping=getattr(base_cfg, "final_logit_softcapping", None),
+            # So that DFlash's conversion warns about them on a fake base too.
+            **{
+                name: getattr(base_cfg, name)
+                for name in _LOGIT_SCALING_FIELDS
+                if getattr(base_cfg, name, None) is not None
+            },
         )
         model = cls(config)
         # Load lm_head, embed_tokens, and (for known models) the final norm into the model.
diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py
index 44cc19ed90c..1d3ae5d8af7 100644
--- a/tests/unit/torch/speculative/plugins/test_fakebase.py
+++ b/tests/unit/torch/speculative/plugins/test_fakebase.py
@@ -25,6 +25,7 @@
 import transformers
 
 from modelopt.torch.speculative.plugins.modeling_fakebase import (
+    _LOGIT_SCALING_FIELDS,
     FakeBaseConfig,
     FakeBaseModel,
     find_final_norm_path,
@@ -71,6 +72,20 @@ def test_fakebase_local_happy_path(fake_checkpoint):
     assert model.embed_tokens.weight.shape == torch.Size([_VOCAB_SIZE, _HIDDEN_SIZE])
 
 
+@pytest.mark.parametrize("softcap", [None, 30.0])
+def test_fakebase_carries_the_logit_softcap(fake_checkpoint, fake_config, softcap):
+    """dflash_teacher_softcap reads the cap off the base config, a fake one included."""
+    fake_config.final_logit_softcapping = softcap
+    assert FakeBaseModel.from_source(str(fake_checkpoint)).config.final_logit_softcapping == softcap
+
+
+@pytest.mark.parametrize("field", _LOGIT_SCALING_FIELDS)
+def test_fakebase_carries_the_logit_scaling(fake_checkpoint, fake_config, field):
+    """DFlash's conversion warns about logit scaling off the base config, a fake one included."""
+    setattr(fake_config, field, 8.0)
+    assert getattr(FakeBaseModel.from_source(str(fake_checkpoint)).config, field) == 8.0
+
+
 def test_fakebase_missing_index_raises(tmp_path, fake_config):
     with pytest.raises(FileNotFoundError, match="safetensors"):
         FakeBaseModel.from_source(str(tmp_path))
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index 6dd54cef812..19a70e25249 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -18,9 +18,11 @@
 GPU-dependent tests (training forward, module forward) are in tests/gpu/.
 """
 
+import copy
 import json
 import logging
 import os
+import warnings
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
@@ -34,7 +36,13 @@
     tf_modelopt_state_and_output_tester,
 )
 from torch.nn.attention.flex_attention import create_mask
-from transformers import AutoModelForCausalLM, Gemma4Config, Gemma4ForConditionalGeneration
+from transformers import (
+    AutoModelForCausalLM,
+    Gemma4Config,
+    Gemma4ForConditionalGeneration,
+    GraniteConfig,
+    GraniteForCausalLM,
+)
 from transformers.modeling_layers import GradientCheckpointingLayer
 
 import modelopt.torch.opt as mto
@@ -65,7 +73,8 @@
     MasterWeightAdamW,
     VerifyMasterWeightsCallback,
 )
-from modelopt.torch.speculative.utils import AcceptanceRateValidation
+from modelopt.torch.speculative.plugins.modeling_fakebase import FakeBaseModel
+from modelopt.torch.speculative.utils import AcceptanceRateValidation, load_vlm_or_llm
 from modelopt.torch.utils.plugins.transformers_dataset import LanguageDataCollator
 
 BLOCK_SIZE = 4
@@ -1623,9 +1632,14 @@ class TestTeacherLogits:
     """
 
     @staticmethod
-    def _model():
+    def _model(softcap=None):
+        """A converted tiny model; ``softcap`` sets the base's cap and turns the switch on."""
         model = get_tiny_llama(num_hidden_layers=4)
-        mtsp.convert(model, [("dflash", get_dflash_config())])
+        config = get_dflash_config()
+        if softcap is not None:
+            model.config.final_logit_softcapping = softcap
+            config["dflash_teacher_softcap"] = True
+        mtsp.convert(model, [("dflash", config)])
         return model
 
     @staticmethod
@@ -1668,8 +1682,28 @@ def test_token_ids_pick_entries_of_the_full_rows(self):
             model._teacher_logits(outputs, labels, token_ids), rows.gather(-1, token_ids)
         )
 
-    def test_handed_over_logits_are_gathered_not_recomputed(self):
-        model = self._model()
+    @pytest.mark.parametrize("narrow", [False, True], ids=["full", "token_ids"])
+    def test_softcap_caps_the_projected_teacher(self, narrow):
+        cap = 0.5
+        model = self._model(softcap=cap).float()
+        hidden = 10 * self._hidden(model)
+        labels = self._labels()
+        rows = self._rows(model._base_model_lm_head(hidden), labels - 1)
+        assert rows.abs().max() > cap, "fixture does not reach the cap"
+        want = cap * torch.tanh(rows / cap)
+        token_ids = None
+        if narrow:
+            token_ids = torch.randint(0, model.config.vocab_size, (*labels.shape, 5))
+            want = want.gather(-1, token_ids)
+        got = model._teacher_logits(
+            DFlashBaseModelOutput(None, base_hidden=hidden), labels, token_ids
+        )
+        torch.testing.assert_close(got, want)
+
+    @pytest.mark.parametrize("softcap", [None, 0.5], ids=["uncapped", "capped"])
+    def test_handed_over_logits_are_gathered_not_recomputed(self, softcap):
+        """Handed-over logits are the model's own: gathered as they are, never capped again."""
+        model = self._model(softcap)
         logits = torch.randn(2, SEQ_LEN, model.config.vocab_size)
         outputs = DFlashBaseModelOutput(None, logits=logits)
         labels = self._labels()
@@ -1763,12 +1797,19 @@ def test_online_teacher_is_the_causal_lm_distribution(self):
 _VISION_START, _IMAGE = 28, 30
 
 
-def _tiny_vlm(arch=None):
-    """A tiny Qwen3-VL converted online, in fp32; ``arch`` updates the draft's architecture."""
+def _tiny_vlm(arch=None, softcap=None):
+    """A tiny Qwen3-VL converted online, in fp32; ``arch`` updates the draft's architecture.
+
+    ``softcap`` names a cap in its text config, which Qwen3-VL's own forward never applies, and
+    turns ``dflash_teacher_softcap`` on.
+    """
     model = get_tiny_qwen3vl(num_hidden_layers=4)
     model.config.vision_start_token_id, model.config.image_token_id = _VISION_START, _IMAGE
     config = get_dflash_config()
     config["dflash_architecture_config"].update(arch or {})
+    if softcap is not None:
+        model.config.text_config.final_logit_softcapping = softcap
+        config["dflash_teacher_softcap"] = True
     mtsp.convert(model, [("dflash", config)])
     return model.float()
 
@@ -1896,6 +1937,18 @@ def test_text_and_image_batches_read_the_same_final_hidden_state(self):
         text = model._online_base_outputs(batch["input_ids"], batch["attention_mask"]).base_hidden
         torch.testing.assert_close(image[:, :4], text[:, :4])
 
+    def test_softcap_caps_the_image_batch_teacher(self):
+        cap = 0.1
+        model = _tiny_vlm(softcap=cap)
+        batch = _image_batch()
+        outputs = model._online_base_outputs(**batch)
+        labels = self._labels()
+        rows = self._own_rows(model, batch, labels - 1)
+        assert rows.abs().max() > cap, "fixture does not reach the cap"
+        torch.testing.assert_close(
+            model._teacher_logits(outputs, labels), cap * torch.tanh(rows / cap)
+        )
+
     @pytest.mark.parametrize(
         "arch",
         [
@@ -1980,3 +2033,112 @@ def drop_the_last_row(_module, _args, output):
         with pytest.raises(RuntimeError, match="once over the whole sequence"):
             model._online_base_outputs(**batch)
         assert base._forward_hooks == hooks
+
+
+class TestTeacherSoftcapAtConvert:
+    """``dflash_teacher_softcap`` alone decides the cap, checked against the base at build time."""
+
+    @staticmethod
+    def _base(multimodal=True, cap=30.0):
+        """A tiny base capping its logits at ``cap``; a nested text config marks a multimodal one."""
+        model = get_tiny_llama(num_hidden_layers=4)
+        model.config.num_orig_hidden_layers = 4
+        text_config = model.config
+        if multimodal:
+            text_config = model.config.text_config = copy.deepcopy(model.config)
+        text_config.final_logit_softcapping = cap
+        return model
+
+    def test_a_base_without_a_cap_is_rejected(self):
+        config = get_dflash_config()
+        config["dflash_teacher_softcap"] = True
+        with pytest.raises(ValueError, match="final_logit_softcapping"):
+            mtsp.convert(get_tiny_llama(num_hidden_layers=4), [("dflash", config)])
+
+    @pytest.mark.parametrize("switch", [False, True], ids=["switch-off", "switch-on"])
+    @pytest.mark.parametrize("multimodal", [False, True], ids=["text-only", "multimodal"])
+    def test_a_capping_base_converts_with_the_cap_the_switch_sets(self, multimodal, switch):
+        """Online, with KD reading the teacher: a multimodal base gets the same rule."""
+        model = self._base(multimodal)
+        config = get_dflash_config()
+        config["dflash_teacher_softcap"] = switch
+        mtsp.convert(model, [("dflash", config)])
+        assert model.dflash_self_logit_distillation
+        assert model._teacher_softcap == (30.0 if switch else None)
+
+    @pytest.mark.parametrize(
+        ("saved", "cap"),
+        [(True, 30.0), (False, None), (None, None)],
+        ids=["on", "off", "saved-before-the-switch"],
+    )
+    def test_restore_reads_the_switch_from_the_saved_config(self, saved, cap):
+        """A checkpoint saved before the switch existed restores with it off."""
+        model = self._base()
+        config = get_dflash_config()
+        config["dflash_teacher_softcap"] = bool(saved)
+        mtsp.convert(model, [("dflash", config)])
+        state = mto.modelopt_state(model)
+        if saved is None:
+            del state["modelopt_state_dict"][0][1]["config"]["dflash_teacher_softcap"]
+        restored = mto.restore_from_modelopt_state(self._base(), state)
+        assert restored._teacher_softcap == cap
+
+
+class TestLogitScalingWarning:
+    """Conversion names the logit scaling that the teacher, lm_head's raw projection, leaves out.
+
+    Once, on rank 0, and only when a loss term reads the teacher.
+    """
+
+    @staticmethod
+    def _granite(logits_scaling):
+        """A tiny Granite, whose CausalLM divides its logits by ``logits_scaling``."""
+        config = GraniteConfig(
+            vocab_size=32,
+            hidden_size=32,
+            intermediate_size=32,
+            num_hidden_layers=4,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            max_position_embeddings=64,
+            logits_scaling=logits_scaling,
+        )
+        return GraniteForCausalLM(config)
+
+    @staticmethod
+    def _warnings(model, kd=True, offline=None):
+        config = get_dflash_config(offline=offline)
+        config["dflash_self_logit_distillation"] = kd
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            mtsp.convert(model, [("dflash", config)])
+        return [str(w.message) for w in caught if "teacher does not apply" in str(w.message)]
+
+    def test_a_scaling_base_warns_once(self):
+        (message,) = self._warnings(self._granite(8.0))
+        assert "logits_scaling=8.0" in message
+
+    def test_an_offline_fake_base_warns_too(self, tmp_path):
+        """Offline VLMs always train on a fake base, which carries the scaling fields over."""
+        self._granite(8.0).save_pretrained(tmp_path)
+        model = load_vlm_or_llm(str(tmp_path), use_fake_base=True, use_offline_training=True)
+        assert isinstance(model, FakeBaseModel)
+        (message,) = self._warnings(model, offline=True)
+        assert "logits_scaling=8.0" in message
+
+    @pytest.mark.parametrize(
+        ("base", "kd"),
+        [("llama", True), ("granite-unscaled", True), ("granite", False)],
+        ids=["llama", "identity-scaling", "teacher-unread"],
+    )
+    def test_no_warning_otherwise(self, base, kd):
+        model = {
+            "llama": lambda: get_tiny_llama(num_hidden_layers=4),
+            "granite-unscaled": lambda: self._granite(1.0),
+            "granite": lambda: self._granite(8.0),
+        }[base]()
+        assert not self._warnings(model, kd)
+
+    def test_only_rank_0_warns(self, monkeypatch):
+        monkeypatch.setattr("modelopt.torch.utils.distributed.is_master", lambda group=None: False)
+        assert not self._warnings(self._granite(8.0))

From 6b39cc14212b6181c768184fa7c6c93a9da20c69 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 15:06:04 +0000
Subject: [PATCH 3/4] feat(speculative): compose the DFlash-family objective
 from dflash_loss_terms

Which divergence a DFlash-family draft trains on was tied to its variant: DFlash, DFlash2 and
LiLiCorr picked KD or CE by dflash_self_logit_distillation (DFlash2 also tv or lambda by
dflash_lk_loss_type), DSpark always blended CE and TVD, and Domino always trained CE. The terms
already live in one registry that names no variant, so the objective becomes a config choice.

- DFlashConfig.dflash_loss_terms takes {term: weight} over the registry's names (ce, kd, tvd,
  tv, lambda), summed in the order given; lambda keeps dflash_lk_ce_scale/_decay, whose
  descriptions stop saying DFlash2 only. Unknown names, negative or non-finite weights and an
  all-zero objective fail config validation, as does setting it together with a legacy
  selector moved off its default. The legacy check compares against the field defaults
  rather than model_fields_set, because main.py converts from model_dump() and modelopt_state
  restores from it, and both mark every field as set.
- No DSpark or Domino recipe pins dflash_self_logit_distillation: false any more: not
  dspark.yaml or domino.yaml, the Nemotron-3.5 DSpark warm-start recipe, or the Qwen3-8B
  Domino launcher example. Neither variant reads it: each states its own default objective,
  and only DFlash's default objective, which DFlash2 and LiLiCorr inherit, and DFlash2's
  modify() read the selector. But next to dflash_loss_terms the pin had to be flipped back
  first. The DFlash2 and LiLiCorr recipes read it and keep theirs; the DSpark and Domino
  skill notes now say the two ignore it.
- Unset, each variant keeps its objective exactly: a variant now states its default as
  _default_loss_terms(config), and _build_objective turns either source into the same
  [(name, weight, term)] list it built before.
- Any term works for any variant, online and offline: online every variant now has the
  teacher, and the offline entry check still follows the terms' needs_teacher declarations.
  DSpark declares its confidence head there, whose target reads the TVD, now that its
  objective may contain no TVD term.
- Every variant reports each term as loss_<term>, a term of weight 0 included. Domino built
  its metrics from scratch, harmless while its objective was a fixed CE equal to loss_final;
  it now keeps the objective's terms next to loss_final and loss_backbone.
- DIVERGENCES is keyed by the config's DFlashLossTerm, so mypy checks its names against the
  config, and a test checks that every configurable name has a term.

Variant-owned terms keep their own knobs: the DFlash2 selector, the LiLiCorr lattice, the DSpark
confidence head, and Domino's lambda schedule, which now blends the configured objective on its
corrected and backbone logits.

Tested: tests/unit/torch/speculative and tests/unit/torch/export/test_hf_spec_rope_export.py
583 passed, and tests/unit/recipe 540 passed, on transformers 5.5.4; the new
tests take at most 0.45 s each on 4 CPUs. They cover the config errors, that the unset
objective of each variant is the one it built before and equals the same terms set explicitly
(lambda's bound parameters included), the terms surviving a model_dump() convert and a
modelopt_state restore in order, a composed objective training every variant online with a
gradient on every draft parameter and every term reported, which configs make offline
batches carry the base hidden state, and that every shipped DSpark and Domino recipe takes a
dflash_loss_terms override alone through load_recipe and main.py's model_dump() round trip.
Without one, dspark.yaml and domino.yaml convert to the objective they built before, and a
DSpark or Domino step is bitwise the same with the selector on or off.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 modelopt/torch/speculative/config.py          |  73 ++++++-
 .../speculative/plugins/dflash_losses.py      |  10 +-
 .../torch/speculative/plugins/hf_dflash.py    |  15 +-
 .../torch/speculative/plugins/hf_dflash2.py   |  20 +-
 .../torch/speculative/plugins/hf_domino.py    |  11 +-
 .../torch/speculative/plugins/hf_dspark.py    |  14 +-
 .../general/speculative_decoding/domino.yaml  |   2 -
 .../general/speculative_decoding/dspark.yaml  |   3 -
 .../dspark_warmstart.yaml                     |   1 -
 .../references/algorithms/domino.md           |   2 +-
 .../references/algorithms/dspark.md           |   2 +-
 tests/unit/recipe/test_loader.py              |  30 ++-
 .../speculative/plugins/test_hf_dflash.py     | 186 ++++++++++++++++++
 .../Qwen/Qwen3-8B/hf_online_domino.yaml       |   1 -
 14 files changed, 318 insertions(+), 52 deletions(-)

diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py
index 902f71ef19f..86e4bd77689 100644
--- a/modelopt/torch/speculative/config.py
+++ b/modelopt/torch/speculative/config.py
@@ -16,9 +16,9 @@
 """Configurations for speculative decoding modes."""
 
 from copy import deepcopy
-from typing import Literal
+from typing import Annotated, Literal
 
-from pydantic import model_validator
+from pydantic import Field, model_validator
 
 from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField
 
@@ -75,6 +75,17 @@ def _get_dflash_default_config():
     },
 }
 
+# The terms a DFlash-family block objective is built from; plugins/dflash_losses.py implements them.
+DFlashLossTerm = Literal["ce", "kd", "tvd", "tv", "lambda"]
+
+# The per-projector_type selectors of the block objective that dflash_loss_terms replaces.
+_LEGACY_OBJECTIVE_FIELDS = (
+    "dflash_self_logit_distillation",
+    "dflash_lk_loss_type",
+    "dflash_ce_loss_alpha",
+    "dflash_l1_loss_alpha",
+)
+
 
 class DFlashConfig(ModeloptBaseConfig):
     """DFlash config for block-wise parallel speculative decoding."""
@@ -98,6 +109,30 @@ class DFlashConfig(ModeloptBaseConfig):
         default=True, description="Whether to freeze base model during DFlash module training."
     )
 
+    dflash_loss_terms: (
+        dict[DFlashLossTerm, Annotated[float, Field(ge=0.0, allow_inf_nan=False)]] | None
+    ) = ModeloptField(
+        default=None,
+        description=(
+            "The block objective as ``{term: weight}``, summed in the order given, e.g. "
+            "``{ce: 0.1, tvd: 0.9}``. Terms: ``ce`` (cross-entropy on the labels), ``kd`` "
+            "(cross-entropy against the base model's distribution), ``tvd`` (total variation "
+            "to the base model's distribution), ``tv`` (``1 - q(gold)``, total variation to "
+            "the label) and ``lambda`` (``ce`` and ``tv`` blended by acceptance, see "
+            "dflash_lk_ce_scale). Any term works with any projector_type, online and offline; "
+            "``kd`` and ``tvd`` read the base model's distribution, which offline batches must "
+            "then carry. Weights are >= 0 with at least one > 0; a term of weight 0 is "
+            "reported without being trained on. None (default) keeps each projector_type's "
+            "objective: ``kd`` or ``ce`` per dflash_self_logit_distillation for DFlash and "
+            "LiLiCorr, the same or dflash_lk_loss_type's term for DFlash2, "
+            "dflash_ce_loss_alpha * ``ce`` + dflash_l1_loss_alpha * ``tvd`` for DSpark, and "
+            "``ce`` for Domino. Setting it requires those selectors to keep their defaults. "
+            "The terms a variant adds itself (the DFlash2 selector, the LiLiCorr lattice, the "
+            "DSpark confidence head) and Domino's blend of the objective on its corrected and "
+            "backbone logits keep their own settings."
+        ),
+    )
+
     dflash_self_logit_distillation: bool = ModeloptField(
         default=True, description="Whether to use logit distillation from base model."
     )
@@ -162,7 +197,7 @@ class DFlashConfig(ModeloptBaseConfig):
     dflash_use_torch_compile: bool = ModeloptField(
         default=False,
         description=(
-            "Whether to torch.compile the draft decoder stack and the DSpark TVD chunk in "
+            "Whether to torch.compile the draft decoder stack and the TVD term's chunk in "
             "training. Compiles once, on the first step, and needs a working Inductor "
             "toolchain."
         ),
@@ -438,8 +473,9 @@ class DFlashConfig(ModeloptBaseConfig):
         default=1.0,
         ge=0.0,
         description=(
-            "DFlash2 only: scale of the CE share in the dflash_lk_loss_type='lambda' "
-            "blend. 1.0 starts the run as pure CE. Ignored for other loss types."
+            "Scale of the CE share in the 'lambda' blend (dflash_lk_loss_type='lambda', or the "
+            "'lambda' term of dflash_loss_terms). 1.0 starts the run as pure CE. Ignored "
+            "without that term."
         ),
     )
 
@@ -447,9 +483,8 @@ class DFlashConfig(ModeloptBaseConfig):
         default=1.0,
         ge=0.0,
         description=(
-            "DFlash2 only: how fast the CE share decays as acceptance rises in the "
-            "dflash_lk_loss_type='lambda' blend. 0 pins the blend at dflash_lk_ce_scale. "
-            "Ignored for other loss types."
+            "How fast the CE share decays as acceptance rises in the 'lambda' blend. 0 pins "
+            "the blend at dflash_lk_ce_scale. Ignored without that term."
         ),
     )
 
@@ -480,6 +515,28 @@ def _check_dpace_alpha(self) -> "DFlashConfig":
                 )
         return self
 
+    @model_validator(mode="after")
+    def _check_loss_terms(self) -> "DFlashConfig":
+        if self.dflash_loss_terms is None:
+            return self
+        if not any(self.dflash_loss_terms.values()):
+            raise ValueError(
+                f"dflash_loss_terms needs a term of weight > 0, got {self.dflash_loss_terms}."
+            )
+        # Against the defaults, not model_fields_set: main.py converts from model_dump() and
+        # modelopt_state restores from it, which marks every field as set.
+        moved = [
+            name
+            for name in _LEGACY_OBJECTIVE_FIELDS
+            if getattr(self, name) != type(self).model_fields[name].default
+        ]
+        if moved:
+            raise ValueError(
+                f"dflash_loss_terms sets the block objective, so {moved} must keep their "
+                "defaults; express the objective in dflash_loss_terms alone."
+            )
+        return self
+
 
 class MedusaConfig(ModeloptBaseConfig):
     """Medusa config."""
diff --git a/modelopt/torch/speculative/plugins/dflash_losses.py b/modelopt/torch/speculative/plugins/dflash_losses.py
index 9b71fb1cf42..d18a1aef602 100644
--- a/modelopt/torch/speculative/plugins/dflash_losses.py
+++ b/modelopt/torch/speculative/plugins/dflash_losses.py
@@ -42,9 +42,9 @@
 (: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`.
+The divergences a block objective is built from live in :data:`DIVERGENCES`, under the
+names ``dflash_loss_terms`` takes. An objective is an ordered list of
+``(name, weight, term)`` entries, summed left to right by :func:`objective_loss`.
 """
 
 import functools
@@ -54,6 +54,8 @@
 import torch.nn.functional as F
 import torch.utils.checkpoint
 
+from ..config import DFlashLossTerm
+
 __all__ = [
     "DIVERGENCES",
     "DraftLogits",
@@ -387,7 +389,7 @@ def __call__(self, ctx, source):
 
 
 # The closed set of terms an objective is built from, by the name its metric reports.
-DIVERGENCES = {
+DIVERGENCES: dict[DFlashLossTerm, type[_Divergence]] = {
     "ce": _CrossEntropy,
     "kd": _SoftCrossEntropy,
     "tvd": _TeacherTotalVariation,
diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 02333a77917..3fcf336e00d 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -500,8 +500,19 @@ def _build_draft_module(self, dflash_config):
         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)]
+        """The block objective, ``[(name, weight, term)]``, from ``dflash_loss_terms`` if set."""
+        terms = config.dflash_loss_terms
+        if terms is None:
+            terms = self._default_loss_terms(config)
+        blend = {"scale": config.dflash_lk_ce_scale, "decay": config.dflash_lk_ce_decay}
+        return [
+            objective_term(name, weight, **(blend if name == "lambda" else {}))
+            for name, weight in terms.items()
+        ]
+
+    def _default_loss_terms(self, config):
+        """The variant's ``{term: weight}`` without ``dflash_loss_terms``: KD from the base, or CE."""
+        return {"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."""
diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py
index f716f4eaad1..8d5b047da1d 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:`_build_objective` and :meth:`_compute_loss`.
+:meth:`_default_loss_terms` 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, objective_term
+from .dflash_losses import LossOutput, _block_tokens
 from .hf_dflash import HFDFlashModel
 from .modeling_dflash2 import DFlash2Module
 
@@ -145,17 +145,11 @@ 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 _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 _default_loss_terms(self, config):
+        """DFlash's terms, or the label TV / CE-TV blend ``dflash_lk_loss_type`` selects."""
+        if config.dflash_lk_loss_type == "ce":
+            return super()._default_loss_terms(config)
+        return {config.dflash_lk_loss_type: 1.0}
 
     def _compute_loss(self, ctx):
         """Backbone DFlash loss plus the candidate-selector cross-entropy.
diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py
index 2dd4cecfd94..b33734728f4 100644
--- a/modelopt/torch/speculative/plugins/hf_domino.py
+++ b/modelopt/torch/speculative/plugins/hf_domino.py
@@ -51,7 +51,8 @@
 
     loss = (1 - lambda_base) * final_loss + lambda_base * base_loss
 
-where ``final_loss`` is CE on the corrected logits and ``base_loss`` is CE on the
+where ``final_loss`` is the block objective (CE unless ``dflash_loss_terms`` sets
+another) on the corrected logits and ``base_loss`` the same objective on the
 backbone-only logits. ``lambda_base`` decays linearly from ``lambda_base_start``
 to 0 over ``lambda_base_decay_ratio`` of training (curriculum: learn a good
 parallel backbone first, then the causal correction). The schedule is driven by
@@ -65,7 +66,7 @@
 from transformers.trainer_callback import TrainerCallback
 
 from ..dflash.conversion import DominoDMRegistry
-from .dflash_losses import DraftLogits, LossOutput, _block_tokens, objective_loss, objective_term
+from .dflash_losses import DraftLogits, LossOutput, _block_tokens, objective_loss
 from .hf_dflash import HFDFlashModel
 from .modeling_domino import DominoModule
 
@@ -133,9 +134,9 @@ def get_exporter(self):
     def _current_lambda_base(self) -> float:
         return float(getattr(self, "_lambda_base", self.dflash_lambda_base_start))
 
-    def _build_objective(self, config):
+    def _default_loss_terms(self, config):
         """CE on the labels, for the corrected logits and for the backbone's."""
-        return [objective_term("ce", 1.0)]
+        return {"ce": 1.0}
 
     @property
     def _position_weighting(self):
@@ -178,9 +179,9 @@ def _compute_loss(self, ctx):
             backbone_loss, _ = objective_loss(self._objective, ctx, "backbone")
             loss = (1.0 - lambda_base) * out.loss + lambda_base * backbone_loss
         metrics = {
+            **out.metrics,
             "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),
         }
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index a4dfbe3032e..4f249dfded5 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -57,6 +57,7 @@
 where ``TVD`` is the total-variation distance between the corrected draft
 distribution and the target (base-model) distribution at the aligned position,
 and the confidence target is ``c* = 1 - 0.5 * TVD`` (the analytical accept rate).
+``dflash_loss_terms`` replaces the first two terms with any composition.
 """
 
 import logging
@@ -66,7 +67,7 @@
 from transformers import PreTrainedModel
 
 from ..dflash.conversion import DSparkDMRegistry
-from .dflash_losses import DraftLogits, LossOutput, _block_tokens, objective_term
+from .dflash_losses import DraftLogits, LossOutput, _block_tokens
 from .hf_dflash import HFDFlashModel
 from .modeling_dspark import DSparkModule
 
@@ -111,12 +112,13 @@ def get_exporter(self):
 
         return DSparkExporter(self)
 
-    def _build_objective(self, config):
+    def _default_loss_terms(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),
-        ]
+        return {"ce": config.dflash_ce_loss_alpha, "tvd": config.dflash_l1_loss_alpha}
+
+    def _reads_teacher(self, config) -> bool:
+        # The confidence head's target, the accept rate 1 - TVD / 2, reads the base distribution.
+        return super()._reads_teacher(config) or config.dflash_confidence_head_alpha > 0
 
     @property
     def _position_weighting(self):
diff --git a/modelopt_recipes/general/speculative_decoding/domino.yaml b/modelopt_recipes/general/speculative_decoding/domino.yaml
index e7237b11b8f..659359be418 100644
--- a/modelopt_recipes/general/speculative_decoding/domino.yaml
+++ b/modelopt_recipes/general/speculative_decoding/domino.yaml
@@ -62,8 +62,6 @@ dflash:
   dflash_block_size: 16
   dflash_num_anchors: 256
   dflash_use_torch_compile: false
-  # Domino does not use target-logit KD; it trains its own base/final CE losses.
-  dflash_self_logit_distillation: false
   # gamma for exponential loss decay (block_size=16 -> 7).
   dflash_loss_decay_factor: 7.0
   # Qwen3 has no native mask token; 151669 is an unused id used by the reference.
diff --git a/modelopt_recipes/general/speculative_decoding/dspark.yaml b/modelopt_recipes/general/speculative_decoding/dspark.yaml
index db69f3c4b9e..fc0cf3e3788 100644
--- a/modelopt_recipes/general/speculative_decoding/dspark.yaml
+++ b/modelopt_recipes/general/speculative_decoding/dspark.yaml
@@ -65,9 +65,6 @@ dflash:
   dflash_block_size: 16
   dflash_num_anchors: 256
   dflash_use_torch_compile: false
-  # DSpark computes the target distribution internally for its TVD/confidence
-  # terms; the DFlash KD path is unused.
-  dflash_self_logit_distillation: false
   # gamma for exponential loss decay (block_size=16 -> 7).
   dflash_loss_decay_factor: 7.0
   # Qwen3 has no native mask token; 151669 is an unused id used by the reference.
diff --git a/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml b/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml
index d96780e05d2..263537a40ed 100644
--- a/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml
+++ b/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml
@@ -81,7 +81,6 @@ dflash:
   # --- training knobs ---
   dflash_num_anchors: 64
   dflash_use_torch_compile: false
-  dflash_self_logit_distillation: false
   dflash_loss_objective: dpace
   dflash_ce_loss_alpha: 0.1
   dflash_l1_loss_alpha: 0.9
diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md
index 8494f6f4b76..f6181b398f4 100644
--- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md
+++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md
@@ -49,7 +49,7 @@ Everything in `dflash.md` applies. Domino adds:
 | `dflash.dflash_lambda_base_start` | 1.0 | Curriculum start weight on the base loss |
 | `dflash.dflash_lambda_base_decay_ratio` | 1.0 | Fraction of training over which `lambda_base` decays to 0 |
 
-`dflash_self_logit_distillation` is **false** — Domino trains its own base/final CE
+Domino ignores `dflash_self_logit_distillation` — it trains its own base/final CE
 losses rather than distilling target logits. Recipe defaults also differ from
 DFlash's: `block_size` 16, `num_anchors` 256, `num_train_epochs` 6,
 `training_seq_len` 3072, `warmup_ratio` 0.04, `dflash_loss_decay_factor` 7.0.
diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md
index bb2c6992ac5..d0bad8ee183 100644
--- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md
+++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md
@@ -88,7 +88,7 @@ Everything in `dflash.md` applies. DSpark adds:
 
 Total loss is `ce_alpha*CE + l1_alpha*TVD + conf_alpha*confidence_BCE`.
 
-`dflash_self_logit_distillation` is **false** for DSpark — it computes the target
+DSpark ignores `dflash_self_logit_distillation` — it computes the target
 distribution internally for the TVD and confidence terms, so the DFlash KD path is
 unused. Recipe defaults also differ from DFlash's: `block_size` 16, `num_anchors` 256,
 `num_train_epochs` 6, `training_seq_len` 3072, `warmup_ratio` 0.04.
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index a02128e5e39..f3943f01f2c 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -59,6 +59,7 @@
 )
 from modelopt.torch.quantization.config import QuantizerAttributeConfig, normalize_quant_cfg_list
 from modelopt.torch.quantization.mode import CalibrateModeRegistry, get_modelike_from_algo_cfg
+from modelopt.torch.speculative.config import DFlashConfig
 
 # ---------------------------------------------------------------------------
 # Static YAML fixtures
@@ -329,13 +330,13 @@ def test_import_resolution_honors_huggingface_alias():
 _PTQ_SCHEMA = "modelopt.recipe.config.ModelOptPTQRecipe"
 
 
-def _all_shipped_ptq_recipe_paths():
-    """Every shipped PTQ recipe, discovered from disk rather than a hardcoded list.
+def _all_shipped_recipe_paths(recipe_type):
+    """Every shipped recipe of ``recipe_type``, discovered from disk rather than a hardcoded list.
 
     A recipe says it is PTQ with ``metadata.recipe_type``, with a ``# modelopt-schema:``
     comment naming :class:`ModelOptPTQRecipe`, or -- for a checkpoint alias, which states
     neither -- by delegating to a recipe that does. All three are picked up, so every
-    shipped recipe is swept by the tests below.
+    shipped recipe is swept by the tests below. Other recipe types state theirs the same way.
     """
     root = files("modelopt_recipes")
     paths = []
@@ -351,7 +352,7 @@ def _all_shipped_ptq_recipe_paths():
         # Ask the loader's own resolver rather than re-deriving the rules here: a
         # top-level ``$import`` does not imply PTQ (an alias of a speculative recipe
         # delegates the same way), and this cannot drift from what load_recipe does.
-        if _peek_recipe_type(path) == RecipeType.PTQ:
+        if _peek_recipe_type(path) == recipe_type:
             paths.append(str(rel.with_suffix("")))
     return paths
 
@@ -359,7 +360,7 @@ def _all_shipped_ptq_recipe_paths():
 # Discovered from disk (not hardcoded) so the smoke tests cover every shipped PTQ
 # recipe — general/, model_type/<model_type>/, and models/<org>/<model_id>/ — and
 # never drift as recipes are added, moved, or removed.
-_BUILTIN_PTQ_RECIPES = _all_shipped_ptq_recipe_paths()
+_BUILTIN_PTQ_RECIPES = _all_shipped_recipe_paths(RecipeType.PTQ)
 
 
 def test_ptq_recipes_are_discovered():
@@ -1059,6 +1060,25 @@ def test_load_recipe_dflash_builtin():
     assert recipe.training.training_seq_len == 4096
 
 
+def test_dspark_and_domino_recipes_take_loss_terms_alone():
+    """Neither variant reads ``dflash_self_logit_distillation``, so none of their recipes pins it.
+
+    ``dflash_loss_terms`` needs the old loss selectors at their defaults; a pin there would have
+    to be undone before a user could set the terms.
+    """
+    checked = []
+    for path in _all_shipped_recipe_paths(RecipeType.SPECULATIVE_DFLASH):
+        projector_type = load_recipe(path).dflash.dflash_architecture_config.get("projector_type")
+        if projector_type not in ("dspark", "domino"):
+            continue
+        recipe = load_recipe(path, overrides=["dflash.dflash_loss_terms={ce: 0.5, tvd: 0.5}"])
+        # main.py converts from a model_dump(), which marks every field as set.
+        dumped = DFlashConfig(**recipe.dflash.model_dump())
+        assert dumped.dflash_loss_terms == {"ce": 0.5, "tvd": 0.5}, path
+        checked.append(path)
+    assert len(checked) >= 3, f"too few DSpark/Domino recipes discovered: {checked}"
+
+
 def test_load_recipe_dflash_missing_section_raises(tmp_path):
     """load_recipe raises ValueError when 'dflash' is absent for a SPECULATIVE_DFLASH recipe."""
     bad = tmp_path / "bad.yml"
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index 19a70e25249..107af93409e 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -24,6 +24,7 @@
 import os
 import warnings
 from types import SimpleNamespace
+from typing import get_args
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -49,7 +50,9 @@
 import modelopt.torch.speculative as mtsp
 import modelopt.torch.speculative.plugins.dflash_flex_attention as dflash_flex_attention
 import modelopt.torch.speculative.plugins.hf_dflash as hf_dflash
+from modelopt.torch.speculative.config import DFlashConfig, DFlashLossTerm
 from modelopt.torch.speculative.plugins.dflash_losses import (
+    DIVERGENCES,
     DraftLogits,
     _block_position_weights,
     _block_targets,
@@ -369,6 +372,189 @@ def test_each_term_reduces_as_before(self, name):
         assert torch.equal(loss, legacy)
 
 
+# What each variant's conversion requires beyond get_dflash_config(), sized for the tiny
+# model: its dflash_architecture_config fields, then its other config fields.
+_VARIANTS = {
+    "dflash": ({}, {}),
+    "dflash2": (
+        {"conv_kernel_size": 2, "conv_group_size": 4, "selector_rank": 8, "selector_top_k": 5},
+        {},
+    ),
+    "lilicorr": (
+        {
+            "lilicorr_candidate_topk": 4,
+            "lilicorr_num_layers": 1,
+            "lilicorr_num_heads": 4,
+            "lilicorr_mlp_ratio": 2.0,
+            "lilicorr_factor_dim": 8,
+            "lilicorr_vector_eps": 1e-4,
+            "lilicorr_logit_scale": 8.0,
+        },
+        {
+            "dflash_lilicorr_w_ce": 0.25,
+            "dflash_lilicorr_w_margin": 0.0,
+            "dflash_lilicorr_w_pen": 0.25,
+        },
+    ),
+    "dspark": ({"markov_rank": 16}, {}),
+    # lambda_base below 1, so that Domino's corrected logits train too.
+    "domino": ({"emb_dim": 16, "gru_hidden_dim": 32}, {"dflash_lambda_base_start": 0.5}),
+}
+
+
+def _variant_model(projector_type, offline=False, arch=None, **overrides):
+    """A tiny model converted to ``projector_type``, with ``arch`` and ``overrides`` applied."""
+    arch_fields, fields = _VARIANTS[projector_type]
+    config = get_dflash_config(offline=offline)
+    config["dflash_architecture_config"].update(
+        projector_type=projector_type, **arch_fields, **(arch or {})
+    )
+    config.update(fields, **overrides)
+    model = get_tiny_llama(num_hidden_layers=4)
+    model.config.num_orig_hidden_layers = 4
+    mtsp.convert(model, [("dflash", config)])
+    return model
+
+
+def _objective_spec(model):
+    """The objective as data: each term's name, weight, class and bound parameters."""
+    return [(name, weight, type(term), vars(term)) for name, weight, term in model._objective]
+
+
+class TestLossTerms:
+    """``dflash_loss_terms`` composes the block objective of any variant from the registry."""
+
+    def test_every_configurable_term_is_implemented(self):
+        assert set(get_args(DFlashLossTerm)) == set(DIVERGENCES)
+
+    @pytest.mark.parametrize(
+        ("terms", "match"),
+        [
+            ({"ce": 1.0, "l1": 1.0}, "'ce', 'kd', 'tvd', 'tv' or 'lambda'"),
+            ({"ce": -0.5}, "greater than or equal to 0"),
+            ({"ce": 0.0}, "weight > 0"),
+        ],
+        ids=["unknown", "negative", "all-zero"],
+    )
+    def test_bad_terms_are_rejected(self, terms, match):
+        with pytest.raises(ValueError, match=match):
+            DFlashConfig(dflash_loss_terms=terms)
+
+    @pytest.mark.parametrize(
+        ("selector", "value"),
+        [
+            ("dflash_self_logit_distillation", False),
+            ("dflash_lk_loss_type", "tv"),
+            ("dflash_ce_loss_alpha", 0.5),
+            ("dflash_l1_loss_alpha", 0.5),
+        ],
+    )
+    def test_terms_and_a_legacy_selector_are_rejected(self, selector, value):
+        """One source for the objective."""
+        with pytest.raises(ValueError, match=selector):
+            DFlashConfig(dflash_loss_terms={"ce": 1.0}, **{selector: value})
+
+    @pytest.mark.parametrize(
+        ("projector_type", "legacy", "terms"),
+        [
+            ("dflash", {}, {"kd": 1.0}),
+            ("dflash", {"dflash_self_logit_distillation": False}, {"ce": 1.0}),
+            (
+                "dflash2",
+                {"dflash_self_logit_distillation": False, "dflash_lk_loss_type": "tv"},
+                {"tv": 1.0},
+            ),
+            (
+                "dflash2",
+                {"dflash_self_logit_distillation": False, "dflash_lk_loss_type": "lambda"},
+                {"lambda": 1.0},
+            ),
+            ("lilicorr", {}, {"kd": 1.0}),
+            ("dspark", {"dflash_ce_loss_alpha": 0.3}, {"ce": 0.3, "tvd": 0.9}),
+            ("domino", {}, {"ce": 1.0}),
+        ],
+        ids=[
+            "dflash-kd",
+            "dflash-ce",
+            "dflash2-tv",
+            "dflash2-lambda",
+            "lilicorr",
+            "dspark",
+            "domino",
+        ],
+    )
+    def test_unset_terms_keep_each_variants_objective(self, projector_type, legacy, terms):
+        """Unset, each variant builds its own objective, which the same terms reproduce."""
+        blend = {"dflash_lk_ce_scale": 0.8, "dflash_lk_ce_decay": 2.0}
+        model = _variant_model(projector_type, **legacy, **blend)
+        assert [(name, weight) for name, weight, _ in model._objective] == list(terms.items())
+        explicit = _variant_model(projector_type, dflash_loss_terms=terms, **blend)
+        assert _objective_spec(explicit) == _objective_spec(model)
+
+    def test_terms_survive_the_config_round_trips(self):
+        """main.py converts from a model_dump() and a restore reads one back: all fields set."""
+        config = DFlashConfig(**get_dflash_config(), dflash_loss_terms={"tvd": 0.5, "ce": 0.5})
+        model = get_tiny_llama(num_hidden_layers=4)
+        mtsp.convert(model, [("dflash", config.model_dump())])
+        restored = mto.restore_from_modelopt_state(
+            get_tiny_llama(num_hidden_layers=4), mto.modelopt_state(model)
+        )
+        assert [name for name, _, _ in restored._objective] == ["tvd", "ce"]
+        assert _objective_spec(restored) == _objective_spec(model)
+
+    @pytest.mark.parametrize(
+        ("projector_type", "terms"),
+        [
+            ("dflash", {"ce": 0.5, "tvd": 0.5}),
+            ("dflash2", {"tv": 0.5, "kd": 0.5}),
+            ("lilicorr", {"tvd": 1.0}),
+            ("dspark", {"kd": 1.0}),
+            ("domino", {"kd": 0.5, "lambda": 0.5, "tvd": 0.0}),
+        ],
+        ids=["dflash", "dflash2", "lilicorr", "dspark", "domino"],
+    )
+    def test_a_composed_objective_trains_every_variant(self, projector_type, terms):
+        """Online, the teacher every variant now has; every draft parameter gets a gradient.
+
+        Every term is reported, one of weight 0 too.
+        """
+        model = _variant_model(projector_type, dflash_loss_terms=terms).float()
+        model.train()
+        out = model(**_dflash_batch(model.config.vocab_size))
+        assert torch.isfinite(out.loss)
+        assert {f"loss_{name}" for name in terms} <= out.draft_metrics.keys()
+        out.loss.backward()
+        ungraded = [
+            name
+            for name, p in model.dflash_module.named_parameters()
+            if p.requires_grad and (p.grad is None or not torch.isfinite(p.grad).all())
+        ]
+        assert not ungraded, ungraded
+
+    @pytest.mark.parametrize(
+        ("projector_type", "terms", "confidence", "needs"),
+        [
+            ("dspark", {"ce": 1.0}, 0.0, False),
+            ("dspark", {"ce": 1.0}, 0.5, True),
+            ("domino", None, 0.0, False),
+            ("domino", {"kd": 1.0}, 0.0, True),
+        ],
+        ids=["dspark-ce", "dspark-ce-confidence", "domino", "domino-kd"],
+    )
+    def test_offline_needs_the_teacher_only_for_terms_that_read_it(
+        self, projector_type, terms, confidence, needs
+    ):
+        """Objective terms and variant terms declare it; nothing else does."""
+        model = _variant_model(
+            projector_type,
+            offline=True,
+            arch={"use_confidence_head": bool(confidence)},
+            dflash_loss_terms=terms,
+            dflash_confidence_head_alpha=confidence,
+        )
+        assert model._needs_base_logits is needs
+
+
 class TestDFlashSaveRestore:
     """Test DFlash model save and restore."""
 
diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml
index 7adf3400cd7..9153dcccacd 100644
--- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml
+++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml
@@ -65,7 +65,6 @@ pipeline:
       - dflash.dflash_num_anchors=512
       - dflash.dflash_loss_decay_factor=7
       - dflash.dflash_mask_token_id=151669
-      - dflash.dflash_self_logit_distillation=false
       - dflash.dflash_lambda_base_start=1.0
       - dflash.dflash_lambda_base_decay_ratio=1.0
     environment:

From 9ae3480f7aa2b7c5dba28f48fa73f8314eb54568 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Mon, 12 Oct 2026 15:09:27 +0000
Subject: [PATCH 4/4] docs(speculative): list dflash_loss_terms and
 dflash_teacher_softcap

CHANGELOG: one New Features (Speculative Decoding) entry per field, two Backward Breaking
Changes entries for what the one teacher changes for existing users, and two Bug Fixes entries.
- Online DFlash, DFlash2 and LiLiCorr now distill text and multimodal batches alike from
  lm_head on the base's final hidden state, which leaves out the logit scaling of Granite,
  Cohere, Falcon-H1 and others; conversion warns about it.
- On a base that caps its logits, Gemma 3n and Gemma 4 VLMs included, the teacher is capped
  only with dflash_teacher_softcap set. A run resumed from an HF-format checkpoint saved
  before this release restores that checkpoint's config, in which the switch is off; one
  resumed from an FSDP-sharded checkpoint converts from the recipe and takes its switch.
- The bug fixes are what the one online forward brings. DSpark's and Domino's image and video
  batches now reach the base model's multimodal forward, where since 0.46 they ran its
  language model on the raw token ids. Like DFlash's, their Gemma 3 and PaliGemma image
  batches reach it without token_type_ids, so on transformers 5.5 it raises and on 5.12 and
  later it attends causally over the image. Every variant's Gemma 4 image and video batches,
  which failed in the vision tower, now pass it their patch positions.

The DFlash recipe table and the agent skill's knob table gain both fields, and the skill's
per-model table says when a Gemma-family base needs the switch, DFlash's default KD included.

No shipped recipe or launcher YAML trains a Gemma-family base (the only Gemma launcher example
benchmarks vLLM MTP), so none sets the switch.

Tested: pre-commit on the three files. On a tiny Qwen3-VL (transformers 5.5.4, 5.12.1 and
5.18.0), a tiny Gemma 3 VLM (5.12.1 and 5.18.0) and a tiny Gemma 4 VLM (5.5.4, 5.12.1 and
5.18.0), two different images give online DSpark and Domino the same target hidden states on
the base branch and different ones here, as they do for DFlash. The image batches of a tiny
Gemma 3 and a tiny PaliGemma VLM, with the token_type_ids their processors return, raise for
every variant on transformers 5.5.4; on 5.12.1 and 5.18.0 they train, and the final hidden
state is bitwise the model's own without those token_type_ids. On the base branch DSpark
and Domino trained on them at 5.5.4 without the image. Resumed through main.py, a checkpoint
that the base branch saved from online DFlash KD on a tiny capping Gemma2 continues uncapped
from its HF-format copy whatever the recipe says, and takes the recipe's switch from its
FSDP-sharded copy (transformers 5.5.4, 5.12.1 and 5.18.0).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 CHANGELOG.rst                                               | 6 ++++++
 examples/speculative_decoding/doc/dflash.md                 | 2 ++
 .../speculative-decoding/references/algorithms/dflash.md    | 2 ++
 3 files changed, 10 insertions(+)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 2907f626226..3a12f4e5b7c 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -18,6 +18,8 @@ Changelog
 *Speculative Decoding*
 
 - Add the **external draft** speculative-decoding mode: train a pretrained standalone draft model against a cached base model, instead of grafting a draft head onto the target. Select it with ``mtsp.convert(draft, [("external", cfg)])`` or the ``speculative_external`` recipe; the draft keeps its own embeddings and lm_head and exports as an ordinary HF checkpoint.
+- Add ``dflash_loss_terms`` to compose the DFlash-family block objective from weighted loss terms (``ce``, ``kd``, ``tvd``, ``tv`` and ``lambda``), e.g. ``{ce: 0.1, tvd: 0.9}``, for any ``projector_type``, online or offline. Left unset, each variant keeps its own objective.
+- Add ``dflash_teacher_softcap`` to cap the distillation teacher with the base model's ``final_logit_softcapping`` (Gemma family), the way the model caps its own logits.
 
 *Megatron Framework (M-LM / M-Bridge)*
 
@@ -37,6 +39,8 @@ 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, 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``.
+- Online DFlash, DFlash2 and LiLiCorr training now distills text and multimodal batches alike from the base model's ``lm_head`` applied to its final hidden state, as offline training and DSpark already did. The teacher therefore leaves out the logit scaling some bases apply at the ``lm_head`` (``logits_scaling``, ``logit_scale``, ``lm_head_multiplier``: Granite, Cohere, Falcon-H1 and others), which conversion now warns about.
+- On a base that caps its logits with ``final_logit_softcapping`` (Gemma family, including the Gemma 3n and Gemma 4 VLMs), online DFlash, DFlash2 and LiLiCorr training now needs ``dflash_teacher_softcap: true`` to keep a capped teacher. A run resumed from an HF-format checkpoint saved before this release restores that checkpoint's config, in which the switch is off, and so continues uncapped; a run resumed from an FSDP-sharded checkpoint takes the switch from the recipe.
 
 **Deprecations**
 
@@ -55,6 +59,8 @@ Changelog
 - Fix ONNX remote AutoQDQ safety benchmarks returning infinite latency by running the generated engine with ``trtexec_safe`` on the configured target.
 - Fix ``INT8_DEFAULT_CFG``, ``FP8_DEFAULT_CFG``, ``NVFP4_DEFAULT_CFG`` and the other presets and recipes that import ``configs/ptq/units/default_disabled_quantizers`` and enable ``*input_quantizer``: since 0.44 they also quantized ``nn.LayerNorm`` inputs, which adds quantization error to the residual stream of models built on ``nn.LayerNorm`` such as BERT-style encoders. Re-quantize affected models; a recipe that wants ``nn.LayerNorm`` quantizers must enable them after importing that unit.
 - Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '<tensor>_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT.
+- Fix online DSpark and Domino training on multimodal batches, which ran the base model's language model on the raw token ids and so dropped the image and video features; they now run the base model's multimodal forward, as DFlash's do. Like DFlash's, their Gemma 3 and PaliGemma image batches reach that forward without ``token_type_ids``, so on transformers 5.5 it raises and on 5.12 and later it attends causally over the image.
+- Fix online DFlash-family training on Gemma 4 image and video batches, which failed in the vision tower because their ``image_position_ids`` and ``video_position_ids`` were not passed on.
 
 0.48.0 (2026-10-xx)
 ^^^^^^^^^^^^^^^^^^^
diff --git a/examples/speculative_decoding/doc/dflash.md b/examples/speculative_decoding/doc/dflash.md
index c468c7652be..ffafdf4a0ae 100644
--- a/examples/speculative_decoding/doc/dflash.md
+++ b/examples/speculative_decoding/doc/dflash.md
@@ -165,6 +165,8 @@ See [`modelopt_recipes/general/speculative_decoding/dflash.yaml`](../../../model
 | `dflash.dflash_loss_objective` | `dpace` | Position weighting: `decay` (static) or `dpace` (dynamic, see below) |
 | `dflash.dflash_dpace_alpha` | 0.5 | D-PACE smoothing factor in (0, 1]; only used when objective is `dpace` |
 | `dflash.dflash_self_logit_distillation` | true | Use target model logits as soft labels (vs hard CE) |
+| `dflash.dflash_loss_terms` | unset | Block objective as `{term: weight}` over `ce`, `kd`, `tvd`, `tv` and `lambda`, e.g. `{ce: 0.1, tvd: 0.9}`; unset keeps the variant's own |
+| `dflash.dflash_teacher_softcap` | false | Cap the teacher with the base's `final_logit_softcapping` (Gemma family) |
 | `dflash.dflash_mask_token_id` | auto | Token ID for masked positions (see note below) |
 | `dflash.dflash_architecture_config.num_hidden_layers` | 5 | Draft decoder layers |
 | `training.answer_only_loss` | false | Mask loss on non-assistant tokens |
diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md
index 90d7f457369..9450ad868de 100644
--- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md
+++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md
@@ -87,6 +87,7 @@ in `examples/speculative_decoding/README.md#dflash-block-diffusion-for-speculati
 | `dflash.dflash_num_anchors` | 512 | Random anchor positions sampled per sequence |
 | `dflash.dflash_loss_decay_factor` | 4.0 | Exponential decay gamma; 0 disables |
 | `dflash.dflash_self_logit_distillation` | true | Logit distillation from the target |
+| `dflash.dflash_loss_terms` | unset | `{term: weight}` over `ce`/`kd`/`tvd`/`tv`/`lambda`, any variant; unset keeps the variant's objective. Leave the variant's own loss selectors at their defaults with it |
 | `dflash.dflash_architecture_config.num_hidden_layers` | 5 | Draft decoder layers — must equal the dump's draft depth (`--num-draft-layers` on vLLM; hardcoded 5 on HF / TRT-LLM) |
 | `dflash.dflash_mask_token_id` | auto | See *Per-model adjustments* |
 | `dflash.dflash_swa_window_size` | unset | Sliding-window attention for the draft; must be >= `dflash_block_size` |
@@ -108,6 +109,7 @@ Export is automatic: after training, rank 0 exports every `checkpoint-<step>` to
 | `trust_remote_code` MoE with an older transformers pin | Set `OVERRIDE_TRANSFORMERS` in the task environment (MiniMax-M2.7 needs 4.57.1). Set `ACCELERATE_CONFIG` when the model needs FSDP2 via accelerate config rather than transformers-native `ParallelismConfig`. |
 | Very large MoE base | Use the offline variant with `model.use_fake_base_for_offline=true`; plain DDP suffices, so no FSDP2 patches. Set `MIXED_PRECISION: "no"` with `training.bf16=false` if the model requires it. |
 | Draft trained at short context, served long | Set `dflash.dflash_export_rope_scaling` (YaRN); factor = target context / `training_seq_len`. |
+| Base with `final_logit_softcapping` (Gemma family) and a loss term reading the target's distribution (`kd`, `tvd`; DFlash, DSpark and LiLiCorr by default) | Set `dflash.dflash_teacher_softcap=true` to distill from the capped distribution, text and image batches alike. |
 | Multi-node | Set `NUM_NODES` in the environment; `HEAD_NODE_IP` is auto-detected from Slurm. |
 
 ## Success markers