diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 4f71f5c4c2b..b811dddfd6a 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -30,6 +30,7 @@ Changelog
 - 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.
 - With ``training.dataloader_num_workers`` > 0, ``examples/speculative_decoding/main.py`` now yields each training batch as soon as a worker finishes it, so one slow sample, such as a slow hidden-state fetch in streaming training, no longer stalls every DDP rank. Sample order within the prefetch window is then no longer deterministic; set ``training.dataloader_in_order=true`` to restore it (context-parallel runs always keep it).
 - The ``general/speculative_decoding/dflash2`` recipe now enables ``dflash_use_torch_compile`` and ``dflash_use_flex_attention``. With PyTorch 2.9 or 2.10 its ``tf32: true`` makes Inductor fail on the first step; set ``training.tf32=false`` or ``dflash.dflash_use_torch_compile=false`` there.
+- Add ``dflash_fused_loss`` to the DFlash config. DFlash, DFlash2 and LiLiCorr training then computes its loss in row chunks fused with the frozen ``lm_head`` projection instead of materializing the draft's full-vocabulary logits, which saves time and memory on large vocabularies.
 
 *Megatron Framework (M-LM / M-Bridge)*
 
diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py
index 5e1c3ae0794..99fb09148e2 100644
--- a/modelopt/torch/speculative/config.py
+++ b/modelopt/torch/speculative/config.py
@@ -162,6 +162,18 @@ class DFlashConfig(ModeloptBaseConfig):
         ),
     )
 
+    dflash_fused_loss: bool = ModeloptField(
+        default=False,
+        description=(
+            "Compute the training loss from the draft's hidden states in row chunks, without "
+            "materializing its [rows, vocab] logits: each chunk is projected through lm_head "
+            "and reduced to the per-row statistics the loss reads, together with their "
+            "gradients. Supported by the dflash, dflash2 and lilicorr projector types with a "
+            "frozen, bias-free nn.Linear lm_head. On CUDA the per-chunk statistics run through "
+            "torch.compile."
+        ),
+    )
+
     dflash_swa_window_size: int | None = ModeloptField(
         default=None,
         description=(
diff --git a/modelopt/torch/speculative/dflash/dflash_model.py b/modelopt/torch/speculative/dflash/dflash_model.py
index e4cb8753668..03c2c58caa7 100644
--- a/modelopt/torch/speculative/dflash/dflash_model.py
+++ b/modelopt/torch/speculative/dflash/dflash_model.py
@@ -53,4 +53,5 @@ def modify(self, config):
         self.dflash_attention_sink = config.dflash_attention_sink
         self.dflash_init_checkpoint = config.dflash_init_checkpoint
         self.dflash_use_flex_attention = config.dflash_use_flex_attention
+        self.dflash_fused_loss = config.dflash_fused_loss
         self.dflash_export_rope_scaling = config.dflash_export_rope_scaling
diff --git a/modelopt/torch/speculative/plugins/dflash_fused_head.py b/modelopt/torch/speculative/plugins/dflash_fused_head.py
new file mode 100644
index 00000000000..505c82669f9
--- /dev/null
+++ b/modelopt/torch/speculative/plugins/dflash_fused_head.py
@@ -0,0 +1,265 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Fused, row-chunked lm_head statistics for DFlash-family training.
+
+The draft's vocab-wide logits are read only through per-row statistics: distillation
+cross-entropy and total variation against a teacher, label cross-entropy, logsumexp, argmax
+and top-k. Each chunk of rows is projected and reduced to those statistics, and the
+hidden-space gradient of every differentiable one is written in the same pass, so no
+``[rows, vocab]`` tensor outlives its chunk or is recomputed in backward.
+"""
+
+import functools
+from collections.abc import Iterable
+from typing import NamedTuple
+
+import torch
+import torch.nn.functional as F
+
+__all__ = ["TERMS", "HeadStats", "fused_head_stats"]
+
+# Per-row statistics that can carry a gradient back to the hidden states.
+TERMS = ("kd_ce", "tv", "hard_ce", "lse")
+
+
+class HeadStats(NamedTuple):
+    """Per-row statistics of the head's logits; a statistic that was not requested is ``None``."""
+
+    kd_ce: torch.Tensor | None  # -sum(p * log q) against the teacher distribution p
+    tv: torch.Tensor | None  # 0.5 * |p - q|_1: one minus the acceptance rate of q against p
+    hard_ce: torch.Tensor | None  # -log q[target]
+    lse: torch.Tensor | None  # logsumexp of the logits
+    top1: torch.Tensor  # argmax
+    top_ids: torch.Tensor | None  # top-k ids, descending
+    top_logits: torch.Tensor | None  # their logits, differentiable
+
+
+def _transform(raw, scale, cap):
+    """``cap * tanh(scale * raw / cap)`` computed in fp32 and returned in ``raw``'s dtype."""
+    if scale is None and cap is None:
+        return raw
+    out = raw.float()
+    if scale is not None:
+        out = out * scale
+    if cap is not None:
+        out = cap * torch.tanh(out / cap)
+    return out.to(raw.dtype)
+
+
+def _chunk_stats(raw, teacher, targets, teacher_is_logits, scale, cap, want, grads, top_k):
+    """Statistics of one chunk of rows, and each differentiable one's gradient w.r.t. ``raw``."""
+    logits = _transform(raw, scale, cap)
+    z = logits.float()
+    lse = torch.logsumexp(z, dim=-1, keepdim=True)
+    log_q = z - lse
+    q = torch.exp(log_q)
+    vals, dz = {}, {}
+    if teacher is not None:
+        zt = (teacher if teacher_is_logits else _transform(teacher, scale, cap)).float()
+        p = torch.softmax(zt, dim=-1)
+        if "kd_ce" in want:
+            vals["kd_ce"] = -(p * log_q).sum(dim=-1)
+            if "kd_ce" in grads:
+                dz["kd_ce"] = q - p
+        if "tv" in want:
+            vals["tv"] = 0.5 * (q - p).abs().sum(dim=-1)
+            if "tv" in grads:
+                sign = torch.sign(q - p)
+                dz["tv"] = 0.5 * q * (sign - (sign * q).sum(dim=-1, keepdim=True))
+    if "hard_ce" in want:
+        vals["hard_ce"] = -log_q.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
+        if "hard_ce" in grads:
+            label = torch.arange(z.shape[-1], device=z.device) == targets.unsqueeze(-1)
+            dz["hard_ce"] = q - label.float()
+    if "lse" in want:
+        vals["lse"] = lse.squeeze(-1)
+        if "lse" in grads:
+            dz["lse"] = q
+
+    # d logits / d raw through the transform; its dtype casts pass gradients through unchanged.
+    factor = 1.0 if scale is None else scale
+    if cap is not None:
+        th = torch.tanh(raw.float() * (factor / cap))
+        factor = factor * (1.0 - th * th)
+    top_ids = top_logits = top_factor = None
+    if top_k:
+        top_logits, top_ids = logits.topk(top_k, dim=-1)
+        if isinstance(factor, torch.Tensor):
+            top_factor = factor.gather(-1, top_ids)
+        elif scale is not None:
+            top_factor = torch.full(top_ids.shape, scale, device=raw.device)
+    jac_inputs = tuple((dz[t] * factor).to(raw.dtype) for t in grads)
+    stats = tuple(vals.get(t) for t in TERMS)
+    return jac_inputs, stats, logits.argmax(dim=-1), top_ids, top_logits, top_factor
+
+
+@functools.cache
+def _compiled_chunk_stats():
+    return torch.compile(_chunk_stats)
+
+
+def _cat(parts):
+    return None if parts[0] is None else torch.cat(parts)
+
+
+class _FusedHead(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        hidden,
+        weight,
+        teacher,
+        targets,
+        teacher_is_logits,
+        scale,
+        cap,
+        want,
+        grads,
+        top_k,
+        chunk,
+    ):
+        stats_fn = _compiled_chunk_stats() if hidden.device.type == "cuda" else _chunk_stats
+        # The head runs in its own dtype and the statistics in fp32, whatever autocast says.
+        with torch.autocast(hidden.device.type, enabled=False):
+            h = hidden.to(weight.dtype)
+            jac = {t: torch.empty_like(h) for t in grads}
+            parts = []
+            for s in range(0, h.shape[0], chunk):
+                e = min(s + chunk, h.shape[0])
+                rows = None
+                if teacher is not None:
+                    rows = teacher[s:e]
+                    if not teacher_is_logits:
+                        rows = F.linear(rows.to(weight.dtype), weight)
+                jac_inputs, *outs = stats_fn(
+                    F.linear(h[s:e], weight),
+                    rows,
+                    None if targets is None else targets[s:e],
+                    teacher_is_logits,
+                    scale,
+                    cap,
+                    want,
+                    grads,
+                    top_k,
+                )
+                for t, dz in zip(grads, jac_inputs):
+                    torch.mm(dz, weight, out=jac[t][s:e])
+                parts.append(outs)
+        stats = [_cat([part[0][i] for part in parts]) for i in range(len(TERMS))]
+        top1, top_ids, top_logits, top_factor = (
+            _cat([part[i] for part in parts]) for i in range(1, 5)
+        )
+        ctx.save_for_backward(*(jac[t] for t in grads), top_ids, top_factor, weight)
+        ctx.set_materialize_grads(False)  # statistics the loss never reads cost nothing in backward
+        ctx.grads, ctx.chunk, ctx.hidden_dtype = grads, chunk, hidden.dtype
+        ctx.mark_non_differentiable(
+            *(v for t, v in zip(TERMS, stats) if v is not None and t not in grads),
+            top1,
+            *(() if top_ids is None else (top_ids,)),
+        )
+        return (*stats, top1, top_ids, top_logits)
+
+    @staticmethod
+    def backward(ctx, *grad_outs):
+        *jacs, top_ids, top_factor, weight = ctx.saved_tensors
+        grad = None
+        for t, jac in zip(ctx.grads, jacs):
+            g = grad_outs[TERMS.index(t)]
+            if g is not None:
+                part = g.float().unsqueeze(-1) * jac.float()
+                grad = part if grad is None else grad + part
+        g_top = grad_outs[-1]
+        if g_top is not None:
+            g_top = g_top.float() if top_factor is None else g_top.float() * top_factor
+            # sum_k g[r, k] * weight[id[r, k]], gathering the head rows one chunk at a time.
+            part = torch.cat(
+                [
+                    torch.bmm(
+                        g_top[s : s + ctx.chunk].unsqueeze(1).to(weight.dtype),
+                        weight[top_ids[s : s + ctx.chunk]],
+                    ).squeeze(1)
+                    for s in range(0, g_top.shape[0], ctx.chunk)
+                ]
+            ).float()
+            grad = part if grad is None else grad + part
+        return (None if grad is None else grad.to(ctx.hidden_dtype), *([None] * 10))
+
+
+def fused_head_stats(
+    hidden: torch.Tensor,
+    weight: torch.Tensor,
+    *,
+    teacher_hidden: torch.Tensor | None = None,
+    teacher_logits: torch.Tensor | None = None,
+    targets: torch.Tensor | None = None,
+    want: Iterable[str] = (),
+    grads: Iterable[str] = (),
+    top_k: int = 0,
+    logit_scale: float | None = None,
+    soft_cap: float | None = None,
+    chunk: int = 2048,
+) -> HeadStats:
+    """Per-row statistics of the head's logits ``hidden @ weight.T``, never materialized whole.
+
+    On CUDA the per-chunk statistics run through ``torch.compile``.
+
+    Args:
+        hidden: ``[R, H]`` rows. The ``grads`` statistics and ``top_logits`` backpropagate into
+            them, and nothing into ``weight``, which must be frozen.
+        weight: ``[V, H]`` head weight.
+        teacher_hidden: ``[R, H]`` teacher rows to project through the same head. ``kd_ce``
+            and ``tv`` need this or ``teacher_logits``.
+        teacher_logits: ``[R, V]`` teacher logits, used as given.
+        targets: ``[R]`` label ids, needed for ``hard_ce``.
+        want: Statistics to return, named as in :data:`TERMS`; ``grads`` are returned too.
+        grads: Statistics to differentiate w.r.t. ``hidden``.
+        top_k: Also return each row's ``top_k`` ids and their logits.
+        logit_scale: Multiplier the head applies to its projection, if any.
+        soft_cap: ``c`` of a ``c * tanh(x / c)`` cap the head applies after ``logit_scale``.
+        chunk: Rows per projection; bounds the vocab-wide temporaries.
+    """
+    want, grads = set(want), set(grads)
+    if unknown := (want | grads) - set(TERMS):
+        raise ValueError(f"Unknown statistics {sorted(unknown)}; choose from {TERMS}.")
+    if weight.requires_grad:
+        raise ValueError("fused_head_stats computes no gradient for the head weight; freeze it.")
+    if teacher_hidden is not None and teacher_logits is not None:
+        raise ValueError("Pass teacher_hidden or teacher_logits, not both.")
+    if top_k > weight.shape[0]:
+        raise ValueError(f"top_k={top_k} exceeds the vocabulary size {weight.shape[0]}.")
+    teacher = teacher_hidden if teacher_logits is None else teacher_logits
+    want |= grads
+    if not (torch.is_grad_enabled() and hidden.requires_grad):
+        grads = set()
+    if teacher is None and want & {"kd_ce", "tv"}:
+        raise ValueError("kd_ce and tv need teacher_hidden or teacher_logits.")
+    if targets is None and "hard_ce" in want:
+        raise ValueError("hard_ce needs targets.")
+    return HeadStats(
+        *_FusedHead.apply(
+            hidden,
+            weight,
+            teacher,
+            targets,
+            teacher_logits is not None,
+            logit_scale,
+            soft_cap,
+            tuple(t for t in TERMS if t in want),
+            tuple(t for t in TERMS if t in grads),
+            top_k,
+            chunk,
+        )
+    )
diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 1e20d5b3099..6c468d96bec 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -85,6 +85,7 @@
 from ..dflash.conversion import DFlashDMRegistry
 from ..dflash.dflash_model import DFlashModel
 from .dflash_flex_attention import build_draft_block_mask
+from .dflash_fused_head import HeadStats, fused_head_stats
 from .modeling_dflash import (  # noqa: F401
     DFlashAttention,
     DFlashBaseModelOutput,
@@ -133,6 +134,11 @@ def _gather_rows(x, positions):
     return torch.gather(x, 1, idx).reshape(*positions.shape, x.size(-1))
 
 
+def _prediction_rows(label_positions):
+    """The base rows whose next-token logits predict the tokens at ``label_positions``."""
+    return (label_positions - 1).clamp(min=0)
+
+
 def _block_positions(anchor_positions, block_size, offset, seq_len):
     """Index ``anchor + offset + k`` of each block position ``k``, clamped onto the sequence.
 
@@ -210,10 +216,15 @@ def _block_position_weights(mask, scheme, first, *, decay_factor, dpace_alpha, t
     return mask
 
 
+def _masked_accuracy(preds, target_ids, mask):
+    """Fraction of ``preds`` equal to ``target_ids`` where ``mask`` is set (0-dim)."""
+    correct = (preds == target_ids) & (mask > 0.5)
+    return correct.sum().float() / (mask.sum() + 1e-6)
+
+
 def _masked_top1(logits, target_ids, mask):
     """Top-1 accuracy of ``logits`` against ``target_ids`` where ``mask`` is set (0-dim)."""
-    correct = (logits.argmax(dim=-1) == target_ids) & (mask > 0.5)
-    return correct.sum().float() / (mask.sum() + 1e-6)
+    return _masked_accuracy(logits.argmax(dim=-1), target_ids, mask)
 
 
 def _draft_mask_mod(seq_len, anchor_positions, block_keep_mask, block_size, window, causal):
@@ -313,13 +324,15 @@ class _DFlashLossTerms(NamedTuple):
     ``ce_per_token`` is ``None`` on the KD path, which never forms a per-position
     cross-entropy. ``weights`` carries the position weighting (decay or D-PACE) and
     normalizes by ``weight_sum``; ``supervised_mask`` is the unweighted mask the
-    reported accuracy uses.
+    reported accuracy uses. ``head_stats`` is the fused pass's statistics of the draft's
+    logits under ``dflash_fused_loss``, and ``None`` otherwise.
     """
 
     ce_per_token: torch.Tensor | None
     weights: torch.Tensor
     weight_sum: torch.Tensor
     supervised_mask: torch.Tensor
+    head_stats: HeadStats | None = None
 
 
 @DFlashDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"})
@@ -330,6 +343,10 @@ class HFDFlashModel(DFlashModel):
     # the anchor's own position onward, so slot 0 (the given anchor) is never supervised.
     _label_shift = 0
 
+    # Set in the body of each class whose `forward` / `_compute_loss` read the draft's logits
+    # only through the fused pass. Checked on the overriding class, so it is not inherited.
+    _supports_fused_loss = True
+
     @property
     def _base_model(self):
         return self.get_submodule(self.base_model_path)
@@ -540,6 +557,8 @@ def modify(self, config):
                 param.requires_grad = False
 
         self._find_base_model_parts()
+        if self.dflash_fused_loss:
+            self._check_fused_loss()
 
         # Factory hook: subclasses (e.g. Domino) override to build an augmented
         # draft module while reusing all of DFlash's modify() setup.
@@ -580,6 +599,35 @@ def _build_draft_module(self, dflash_config):
         """Build the draft module. Subclasses override to use an augmented module."""
         return DFlashModule(dflash_config)
 
+    def _check_fused_loss(self):
+        """Refuse ``dflash_fused_loss`` where it would be ignored or would drop a gradient."""
+        for method in ("forward", "_compute_loss"):
+            owner = next(cls for cls in type(self).__mro__ if method in vars(cls))
+            if not vars(owner).get("_supports_fused_loss", False):
+                raise ValueError(
+                    f"dflash_fused_loss is not supported by {owner.__name__}: its {method} "
+                    "reads the draft's logits directly. Unset dflash_fused_loss."
+                )
+        head = self._base_model_lm_head
+        # The fused pass reads lm_head.weight directly: whatever a bias, a quantizer or a
+        # hook adds in lm_head's forward would be silently skipped.
+        plain_forward = (
+            type(head).forward is torch.nn.Linear.forward
+            and "forward" not in vars(head)
+            and not (head._forward_hooks or head._forward_pre_hooks)
+        )
+        if not plain_forward or head.bias is not None:
+            raise ValueError(
+                "dflash_fused_loss projects with lm_head.weight alone, so it needs a bias-free "
+                f"nn.Linear lm_head with no forward hooks; got {type(head).__name__}. Unset "
+                "dflash_fused_loss."
+            )
+        if head.weight.requires_grad:
+            raise ValueError(
+                "dflash_fused_loss computes no gradient for lm_head, which is trainable here; "
+                "it needs dflash_freeze_base_model=True."
+            )
+
     def _base_device(self):
         """Device the frozen base model sits on. ``meta`` during a from_pretrained restore."""
         if self.dflash_offline:
@@ -922,13 +970,24 @@ def _teacher_logits(self, base_outputs, *, label_positions, token_ids=None):
         projection to just those vocab entries. ``label_positions`` is keyword-only, so every
         call says which positions it passes: handing it the logit rows reads one row early.
         """
-        positions = (label_positions - 1).clamp(min=0)
+        positions = _prediction_rows(label_positions)
         if base_outputs.logits is not None:
             if token_ids is None:
                 return _gather_rows(base_outputs.logits, positions)
             batch = torch.arange(positions.shape[0], device=positions.device)
             batch = batch.view(-1, *[1] * positions.dim())
             return base_outputs.logits[batch, positions.unsqueeze(-1), token_ids]
+        lm_head = self._base_model_lm_head
+        rows = self._teacher_hidden(base_outputs, label_positions=label_positions)
+        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]
+
+    @torch.no_grad()
+    def _teacher_hidden(self, base_outputs, *, label_positions):
+        """The base model's final hidden that predicts the tokens at ``label_positions``: [B, ..., H]."""
         if base_outputs.base_hidden is None:
             raise ValueError(
                 "This objective needs the base model's distribution, but the batch carries "
@@ -937,13 +996,9 @@ def _teacher_logits(self, base_outputs, *, label_positions, token_ids=None):
         lm_head = self._base_model_lm_head
         # A producer can store the hidden states in a wider dtype than the target's weights.
         # Cast before the final norm too, which online training runs in the target's dtype.
+        positions = _prediction_rows(label_positions)
         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]
+        return _apply_base_final_norm(rows, base_outputs.base_hidden_prenorm, self._base_model_norm)
 
     def _compute_loss(
         self,
@@ -960,20 +1015,20 @@ def _compute_loss(
         """Compute weighted cross-entropy (or KD) loss and accuracy.
 
         Args:
-            logits: Draft model output [B, N*block_size, vocab].
+            logits: Draft model output [B, N*block_size, vocab], or ``None`` under
+                ``dflash_fused_loss``, which reads ``draft_hidden`` instead.
             input_ids: Original input token IDs [B, seq_len].
             anchor_positions: Anchor positions per block [B, N].
             block_keep_mask: Valid block mask [B, N].
             loss_mask: Token-level loss mask [B, seq_len].
             draft_hidden: Draft hidden states [B, N*block_size, H] behind ``logits``.
-                Unused here; passed for variants whose head consumes them.
             base_outputs: The step's ``DFlashBaseModelOutput``: the KD teacher under
                 ``dflash_self_logit_distillation``, and read by variants.
             return_terms: Also return the unreduced pieces behind the loss, so a variant
                 can recompose the block objective from a different divergence without
                 rebuilding the target alignment and position weighting.
                 TODO: promote this into a shared divergence seam when the DFlash-family
-                loss code is refactored; DFlash2 is the only consumer today.
+                loss code is refactored; DFlash2 and LiLiCorr are its consumers today.
 
         Returns:
             (loss, accuracy) tuple.
@@ -984,7 +1039,6 @@ def _compute_loss(
         )
         binary_eval_mask = targets.mask.view(-1)
 
-        flat_logits = logits.view(-1, logits.size(-1))
         flat_targets = targets.token_ids.view(-1)
 
         kd = self.dflash_self_logit_distillation
@@ -993,12 +1047,25 @@ def _compute_loss(
                 "dflash_self_logit_distillation distills from base_outputs, but none was "
                 "passed; an override of _compute_loss must forward it."
             )
+        # Without logits (dflash_fused_loss), every statistic of them read below comes from
+        # one fused pass over the draft's hidden states.
+        stats = None
+        if logits is None:
+            stats = self._fused_head_stats(
+                draft_hidden, base_outputs, targets.positions, flat_targets
+            )
+        else:
+            flat_logits = logits.view(-1, logits.size(-1))
         # Non-KD loss is per-token cross-entropy; compute it once (grad enabled) so the
         # D-PACE confidences below can reuse it instead of a second CE pass. The KD path
         # optimizes KL, so its confidences need a dedicated no_grad CE pass.
         loss_per_token = None
         if not kd:
-            loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none")
+            loss_per_token = (
+                stats.hard_ce
+                if stats is not None
+                else F.cross_entropy(flat_logits, flat_targets, reduction="none")
+            )
 
         weight_mask = _block_position_weights(
             targets.mask,
@@ -1009,6 +1076,8 @@ def _compute_loss(
             target_ce=lambda: (
                 loss_per_token.detach()
                 if loss_per_token is not None
+                else stats.hard_ce
+                if stats is not None
                 else F.cross_entropy(flat_logits, flat_targets, reduction="none")
             ).view(targets.mask.shape),
         )
@@ -1017,7 +1086,9 @@ def _compute_loss(
         valid_count = flat_weights.sum() + 1e-6
 
         if valid_count > 1.0:
-            if kd:
+            if kd and stats is not None:
+                loss = (stats.kd_ce * flat_weights).sum() / valid_count
+            elif kd:
                 flat_teacher = self._teacher_logits(base_outputs, label_positions=targets.positions)
                 flat_teacher = flat_teacher.reshape(flat_logits.shape)
                 target_soft = torch.softmax(flat_teacher, dim=-1)
@@ -1028,9 +1099,12 @@ def _compute_loss(
                 loss = (loss_per_token * flat_weights).sum() / valid_count
 
             with torch.no_grad():
-                accuracy = _masked_top1(flat_logits, flat_targets, binary_eval_mask).item()
+                preds = stats.top1 if stats is not None else flat_logits.argmax(dim=-1)
+                accuracy = _masked_accuracy(preds, flat_targets, binary_eval_mask).item()
         else:
-            loss = self._draft_zero_loss(flat_logits.device, through=(flat_logits,))
+            # Under dflash_fused_loss the draft's output is its hidden states: no logits exist.
+            output = draft_hidden if stats is not None else flat_logits
+            loss = self._draft_zero_loss(output.device, through=(output,))
             accuracy = 0.0
 
         if return_terms:
@@ -1039,10 +1113,40 @@ def _compute_loss(
                 weights=flat_weights,
                 weight_sum=valid_count,
                 supervised_mask=binary_eval_mask,
+                head_stats=stats,
             )
             return loss, accuracy, terms
         return loss, accuracy
 
+    def _fused_head_extras(self):
+        """``(statistics with gradient, top_k)`` a variant's own terms read from the fused pass."""
+        return set(), 0
+
+    def _fused_head_stats(self, draft_hidden, base_outputs, label_positions, flat_targets):
+        """Every statistic of the draft's logits the loss reads, from one fused pass."""
+        grads, top_k = self._fused_head_extras()
+        teacher = {}
+        if self.dflash_self_logit_distillation:
+            grads = {*grads, "kd_ce"}
+            if base_outputs.logits is not None:
+                rows = self._teacher_logits(base_outputs, label_positions=label_positions)
+                teacher["teacher_logits"] = rows.flatten(0, -2)
+            else:
+                rows = self._teacher_hidden(base_outputs, label_positions=label_positions)
+                teacher["teacher_hidden"] = rows.flatten(0, -2)
+        else:
+            grads = {*grads, "hard_ce"}
+        return fused_head_stats(
+            draft_hidden.flatten(0, -2),
+            self._base_model_lm_head.weight,
+            targets=flat_targets,
+            # D-PACE weighs positions by the draft's confidence in its label.
+            want={"hard_ce"} if self.dflash_loss_objective == "dpace" else (),
+            grads=grads,
+            top_k=top_k,
+            **teacher,
+        )
+
     def forward(
         self,
         input_ids=None,
@@ -1197,7 +1301,7 @@ def forward(
         )
 
         # 6. Compute loss and accuracy
-        logits = self._base_model_lm_head(hidden)
+        logits = None if self.dflash_fused_loss else self._base_model_lm_head(hidden)
         loss, accuracy = self._compute_loss(
             logits,
             input_ids,
diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py
index 88c1e3e38f7..73ba8e1775e 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash2.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash2.py
@@ -55,6 +55,8 @@ class HFDFlash2Model(HFDFlashModel):
     to it when ``dflash_architecture_config.projector_type == "dflash2"``.
     """
 
+    _supports_fused_loss = True
+
     def _build_draft_module(self, dflash_config):
         """Build the DFlash2 draft module (DFlash backbone + convolutions + selector)."""
         return DFlash2Module(dflash_config)
@@ -108,13 +110,27 @@ def get_exporter(self):
 
         return DFlash2Exporter(self)
 
+    def _fused_head_extras(self):
+        """The selector scores the backbone's top-k candidates."""
+        selector = self.dflash_module.candidate_selector
+        return set(), selector.top_k if self.dflash_selector_loss_alpha > 0 else 0
+
     def _selector_loss(
-        self, logits, target_ids, hidden, predecessor_ids, supervised_mask, loss_weights, weight_sum
+        self,
+        logits,
+        target_ids,
+        hidden,
+        predecessor_ids,
+        supervised_mask,
+        loss_weights,
+        weight_sum,
+        head_stats=None,
     ):
         """Cross-entropy over the selector's candidate set, and its top-1 accuracy.
 
         Args:
-            logits: Backbone logits per block position ``[B, N, block_size, V]``.
+            logits: Backbone logits per block position ``[B, N, block_size, V]``, or ``None``
+                under ``dflash_fused_loss``, whose pass took their top-k into ``head_stats``.
             target_ids: Gold token ids ``[B, N, block_size]``.
             hidden: Backbone hidden states ``[B, N, block_size, H]``.
             predecessor_ids: Teacher-forced predecessor ids ``[B, N, block_size]``.
@@ -131,7 +147,11 @@ def _selector_loss(
         selector = self.dflash_module.candidate_selector
         top_k = selector.top_k
 
-        unary_logits, candidate_ids = logits.topk(top_k, dim=-1)
+        if logits is None:
+            unary_logits = head_stats.top_logits.view(*target_ids.shape, top_k)
+            candidate_ids = head_stats.top_ids.view(*target_ids.shape, top_k)
+        else:
+            unary_logits, candidate_ids = logits.topk(top_k, dim=-1)
 
         # Train on the strict top-k, the candidate set serving actually builds. A gold
         # token the backbone did not propose is a backbone recall failure, not a
@@ -260,13 +280,14 @@ def _compute_loss(
         # The backbone's own mask and position weights: a block position's choice only
         # matters if every earlier one was accepted, which decay / D-PACE model.
         selector_loss, selector_accuracy, selector_coverage = self._selector_loss(
-            logits.reshape(bsz, n_blocks, block_size, -1),
+            None if logits is None else logits.reshape(bsz, n_blocks, block_size, -1),
             target_ids,
             draft_hidden.reshape(bsz, n_blocks, block_size, -1),
             predecessor_ids,
             terms.supervised_mask.view_as(target_ids),
             terms.weights.view_as(target_ids),
             terms.weight_sum,
+            head_stats=terms.head_stats,
         )
         self._selector_metrics = {
             "selector_accuracy": selector_accuracy,
diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py
index b38234d78eb..4d7b604a84a 100644
--- a/modelopt/torch/speculative/plugins/hf_lilicorr.py
+++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py
@@ -102,6 +102,8 @@ class HFLiLiCorrModel(HFDFlashModel):
     to it when ``dflash_architecture_config.projector_type == "lilicorr"``.
     """
 
+    _supports_fused_loss = True
+
     def _build_draft_module(self, dflash_config):
         """Build the LiLiCorr draft module (DFlash backbone + lattice reranker)."""
         return LiLiCorrModule(dflash_config)
@@ -220,13 +222,15 @@ def _compute_loss(
                 "committed target row."
             )
 
-        origin_loss, origin_accuracy = super()._compute_loss(
+        origin_loss, origin_accuracy, terms = super()._compute_loss(
             logits,
             input_ids,
             anchor_positions,
             block_keep_mask,
             loss_mask,
+            draft_hidden=draft_hidden,
             base_outputs=base_outputs,
+            return_terms=True,
         )
         # The backbone's own supervision, before its position weighting, so both losses
         # train the same positions.
@@ -240,6 +244,7 @@ def _compute_loss(
         )
         lilicorr_loss, accuracy, metrics = self._compute_lilicorr_loss(
             logits=logits,
+            head_stats=terms.head_stats,
             target_ids=targets.token_ids,
             # Block position 0 is the anchor -- the one row whose input embedding is a real
             # verified token -- so the lattice covers positions 1..bs-1.
@@ -262,6 +267,10 @@ def _compute_loss(
         self._lilicorr_metrics = metrics
         return origin_loss + lilicorr_loss, accuracy
 
+    def _fused_head_extras(self):
+        """The lattice scores the backbone's top-k candidates against its log-normalizer."""
+        return {"lse"}, self.dflash_module.lilicorr.candidate_topk
+
     def _compute_lilicorr_loss(
         self,
         *,
@@ -272,9 +281,13 @@ def _compute_lilicorr_loss(
         draft_hidden,
         target_hidden,
         base_outputs,
+        head_stats=None,
     ):
         """Score the candidate lattice and evaluate the LiLiCorr objective.
 
+        ``logits`` is ``None`` under ``dflash_fused_loss``, whose pass took the candidates,
+        their logits and the log-normalizer into ``head_stats``.
+
         Returns ``(loss, accuracy, metrics)``. ``loss`` already carries the absolute
         term weights, ``accuracy`` is the greedy path-prefix fraction, and ``metrics``
         holds Python floats read back in a single device sync.
@@ -284,15 +297,22 @@ def _compute_lilicorr_loss(
         bsz, n_blocks, block_size = target_ids.shape
         num_slots = block_size - 1
         topk = head.candidate_topk
-        vocab = logits.shape[-1]
+        vocab = self._base_model_lm_head.weight.shape[0] if logits is None else logits.shape[-1]
         batch_blocks = bsz * n_blocks
 
         if topk > vocab:
             raise ValueError(f"lilicorr_candidate_topk={topk} exceeds the vocabulary size {vocab}.")
 
-        logits_slots = logits.view(bsz, n_blocks, block_size, vocab)[:, :, 1:, :].reshape(
-            batch_blocks, num_slots, vocab
-        )
+        def slots(per_position):
+            """Block positions 1..block_size-1 of a ``[B * N * block_size, ...]`` tensor."""
+            return per_position.view(bsz, n_blocks, block_size, -1)[:, :, 1:].reshape(
+                batch_blocks, num_slots, -1
+            )
+
+        if logits is not None:
+            logits_slots = logits.view(bsz, n_blocks, block_size, vocab)[:, :, 1:, :].reshape(
+                batch_blocks, num_slots, vocab
+            )
         ground_truth = target_ids[:, :, 1:]
         gt_flat = ground_truth.reshape(batch_blocks, num_slots)
         # A block is trained only when every one of its slots is a supervised position,
@@ -304,7 +324,10 @@ def _compute_lilicorr_loss(
             # is the rank order `rank_embedding` is indexed by. Taken on the logits
             # rather than on a log_softmax of them: the map is strictly monotone per row,
             # so the ids are identical, and it avoids a vocab-wide float32 buffer.
-            _, candidate_ids = torch.topk(logits_slots.detach(), k=topk, dim=-1)
+            if logits is None:
+                candidate_ids = slots(head_stats.top_ids)
+            else:
+                _, candidate_ids = torch.topk(logits_slots.detach(), k=topk, dim=-1)
             gt_hits = candidate_ids == gt_flat.unsqueeze(-1)
             gt_in_slot = gt_hits.any(dim=-1)
             gt_indices = gt_hits.long().argmax(dim=-1)
@@ -314,8 +337,12 @@ def _compute_lilicorr_loss(
         # re-attaches the autograd graph, which is how the lattice terms reach the
         # drafter body and not only the head. logsumexp rather than a full log_softmax
         # so backward needs no second vocab-sized buffer.
-        candidate_logits = torch.gather(logits_slots, 2, candidate_ids).float()
-        log_denominator = torch.logsumexp(logits_slots.float(), dim=-1, keepdim=True)
+        if logits is None:
+            candidate_logits = slots(head_stats.top_logits).float()
+            log_denominator = slots(head_stats.lse.unsqueeze(-1))
+        else:
+            candidate_logits = torch.gather(logits_slots, 2, candidate_ids).float()
+            log_denominator = torch.logsumexp(logits_slots.float(), dim=-1, keepdim=True)
         candidate_log_probs = candidate_logits - log_denominator
 
         draft_slots = draft_hidden.view(bsz, n_blocks, block_size, -1)[:, :, 1:, :]
diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py
index 5f2e8adc045..37b4cf70921 100644
--- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py
@@ -395,3 +395,51 @@ def test_backward_produces_finite_grads(self):
         ]
         assert grads, "no draft gradients were produced"
         assert all(torch.isfinite(g).all() for g in grads)
+
+
+class TestDFlashFusedLossGPU:
+    """On CUDA the fused loss runs its statistics through torch.compile; they must still match."""
+
+    DFLASH2_ARCH = {
+        "projector_type": "dflash2",
+        "conv_kernel_size": 2,
+        "conv_group_size": 4,
+        "selector_rank": 8,
+        "selector_top_k": 5,
+    }
+
+    @staticmethod
+    def _step(fused, overrides, arch):
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = get_dflash_config()
+        config.update(dflash_fused_loss=fused, **overrides)
+        config["dflash_architecture_config"].update(arch)
+        mtsp.convert(model, [("dflash", config)])
+        model = model.float().cuda().train()  # fp32, so the two paths differ only by rounding
+        torch.manual_seed(0)
+        input_ids = torch.randint(1, model.config.vocab_size, (2, SEQ_LEN), device="cuda")
+        torch.manual_seed(1)
+        out = model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids))
+        out.loss.backward()
+        params = model.dflash_module.named_parameters()
+        return out, {name: p.grad for name, p in params if p.grad is not None}
+
+    @pytest.mark.parametrize(
+        ("overrides", "arch"),
+        [
+            ({}, {}),  # KD on the base's logits, D-PACE weighting
+            (
+                {"dflash_self_logit_distillation": False, "dflash_lk_loss_type": "lambda"},
+                DFLASH2_ARCH,
+            ),
+        ],
+        ids=["dflash-kd", "dflash2-selector"],
+    )
+    def test_matches_the_materialized_logits(self, overrides, arch):
+        eager, eager_grads = self._step(False, overrides, arch)
+        fused, fused_grads = self._step(True, overrides, arch)
+        torch.testing.assert_close(fused.loss, eager.loss)
+        assert fused.train_acc == eager.train_acc
+        assert fused_grads.keys() == eager_grads.keys()
+        for name, grad in eager_grads.items():
+            torch.testing.assert_close(fused_grads[name], grad, msg=name)
diff --git a/tests/unit/torch/speculative/plugins/test_dflash_fused_head.py b/tests/unit/torch/speculative/plugins/test_dflash_fused_head.py
new file mode 100644
index 00000000000..59d5d247143
--- /dev/null
+++ b/tests/unit/torch/speculative/plugins/test_dflash_fused_head.py
@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU tests for the fused lm_head statistics behind ``dflash_fused_loss``."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from modelopt.torch.speculative.plugins.dflash_fused_head import TERMS, fused_head_stats
+
+ROWS, HIDDEN, VOCAB, TOP_K = 37, 16, 50, 4
+CHUNK = 8  # several chunks and a partial last one
+
+
+def _transform(raw, scale, cap):
+    if scale is not None:
+        raw = raw * scale
+    return raw if cap is None else cap * torch.tanh(raw / cap)
+
+
+@pytest.mark.parametrize("teacher_kind", ["hidden", "logits"])
+@pytest.mark.parametrize(("scale", "cap"), [(None, None), (0.7, None), (0.7, 2.0)])
+def test_matches_autograd_on_the_materialized_logits(teacher_kind, scale, cap):
+    """Every statistic, its gradient under arbitrary per-row upstream weights, and the top-k."""
+    torch.manual_seed(0)
+    weight = torch.randn(VOCAB, HIDDEN)
+    hidden = torch.randn(ROWS, HIDDEN)
+    targets = torch.randint(0, VOCAB, (ROWS,))
+    teacher_hidden = torch.randn(ROWS, HIDDEN)
+    teacher_logits = 3 * torch.randn(ROWS, VOCAB)  # handed over as is, never transformed
+    coefs = {term: torch.randn(ROWS) for term in TERMS}
+    top_coefs = torch.randn(ROWS, TOP_K)
+
+    ref_hidden = hidden.clone().requires_grad_(True)
+    logits = _transform(ref_hidden @ weight.T, scale, cap)
+    if teacher_kind == "hidden":
+        teacher = {"teacher_hidden": teacher_hidden}
+        log_p = _transform(teacher_hidden @ weight.T, scale, cap).log_softmax(-1)
+    else:
+        teacher = {"teacher_logits": teacher_logits}
+        log_p = teacher_logits.log_softmax(-1)
+    log_q = logits.log_softmax(-1)
+    ref = {
+        "kd_ce": -(log_p.exp() * log_q).sum(-1),
+        "tv": 0.5 * (log_q.exp() - log_p.exp()).abs().sum(-1),
+        "hard_ce": F.cross_entropy(logits, targets, reduction="none"),
+        "lse": logits.logsumexp(-1),
+    }
+    top_logits, top_ids = logits.topk(TOP_K, dim=-1)
+    ref_loss = sum((coefs[t] * ref[t]).sum() for t in TERMS) + (top_coefs * top_logits).sum()
+    ref_loss.backward()
+
+    fused_hidden = hidden.clone().requires_grad_(True)
+    kwargs = dict(
+        targets=targets,
+        grads=TERMS,
+        top_k=TOP_K,
+        logit_scale=scale,
+        soft_cap=cap,
+        chunk=CHUNK,
+        **teacher,
+    )
+    stats = fused_head_stats(fused_hidden, weight, **kwargs)
+    loss = sum((coefs[t] * getattr(stats, t)).sum() for t in TERMS)
+    (loss + (top_coefs * stats.top_logits).sum()).backward()
+    with torch.no_grad():  # values without the gradient work, e.g. an eval pass in train mode
+        no_grad_stats = fused_head_stats(fused_hidden, weight, **kwargs)
+
+    for term in TERMS:
+        torch.testing.assert_close(getattr(stats, term), ref[term].detach(), msg=term)
+        torch.testing.assert_close(getattr(no_grad_stats, term), ref[term].detach(), msg=term)
+    assert torch.equal(stats.top1, logits.argmax(-1))
+    assert torch.equal(stats.top_ids, top_ids)
+    torch.testing.assert_close(stats.top_logits, top_logits.detach())
+    torch.testing.assert_close(fused_hidden.grad, ref_hidden.grad)
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index d3bf4162cf7..0aae907c171 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -1727,3 +1727,70 @@ def test_offline_step_never_projects_the_full_sequence(self):
         ).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
+
+
+class TestDFlashFusedLoss:
+    """``dflash_fused_loss`` trains exactly what the loss on materialized logits trains."""
+
+    @staticmethod
+    def _step(fused, offline, **overrides):
+        """One training step's output and draft gradients, from fixed weights, data and anchors."""
+        torch.manual_seed(0)
+        model = get_tiny_llama(num_hidden_layers=4)
+        if offline:
+            model.config.num_orig_hidden_layers = 4
+        config = get_dflash_config(offline=offline)
+        config.update(dflash_fused_loss=fused, **overrides)
+        mtsp.convert(model, [("dflash", config)])
+        model.float().train()  # fp32, so the two paths differ only by rounding
+        batch = _dflash_batch(model.config.vocab_size)
+        if offline:
+            hidden = model.config.hidden_size
+            batch["base_model_outputs"] = {
+                "aux_hidden_states": torch.randn(2, SEQ_LEN, len(model.target_layer_ids) * hidden),
+                "base_model_hidden_states": torch.randn(2, SEQ_LEN, hidden),
+            }
+        torch.manual_seed(1)
+        out = model(**batch)
+        out.loss.backward()
+        params = model.dflash_module.named_parameters()
+        return out, {name: p.grad for name, p in params if p.grad is not None}
+
+    @pytest.mark.parametrize(
+        ("offline", "overrides"),
+        [
+            (False, {}),  # KD on the base's logits, D-PACE weighting
+            (True, {"dflash_loss_objective": "decay", "dflash_loss_decay_factor": 4.0}),
+            (False, {"dflash_self_logit_distillation": False}),  # cross-entropy, D-PACE
+        ],
+        ids=["kd-online", "kd-offline-decay", "ce"],
+    )
+    def test_matches_the_materialized_logits(self, offline, overrides):
+        eager, eager_grads = self._step(False, offline, **overrides)
+        fused, fused_grads = self._step(True, offline, **overrides)
+        torch.testing.assert_close(fused.loss, eager.loss)
+        assert fused.train_acc == eager.train_acc
+        assert fused_grads.keys() == eager_grads.keys()
+        for name, grad in eager_grads.items():
+            torch.testing.assert_close(fused_grads[name], grad, msg=name)
+
+    @pytest.mark.parametrize(
+        ("overrides", "arch", "hook_lm_head"),
+        [
+            ({}, {"projector_type": "dspark", "markov_rank": 4}, False),
+            ({}, {"projector_type": "domino", "gru_hidden_dim": 8, "emb_dim": 8}, False),
+            ({"dflash_freeze_base_model": False}, {}, False),
+            ({}, {}, True),
+        ],
+        ids=["dspark", "domino", "trainable-lm-head", "hooked-lm-head"],
+    )
+    def test_unsupported_setups_are_rejected(self, overrides, arch, hook_lm_head):
+        """Each would make the fused pass skip a term or a gradient the eager loss has."""
+        model = get_tiny_llama(num_hidden_layers=4)
+        if hook_lm_head:
+            model.lm_head.register_forward_hook(lambda module, args, output: output * 2)
+        config = get_dflash_config()
+        config.update(dflash_fused_loss=True, **overrides)
+        config["dflash_architecture_config"].update(arch)
+        with pytest.raises(ValueError, match="dflash_fused_loss"):
+            mtsp.convert(model, [("dflash", config)])
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
index e1ea8c10dc0..67219179bbe 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
@@ -384,6 +384,34 @@ def test_overfits_a_single_batch(self):
         assert out["selector_metrics"]["selector_accuracy"].item() > 0.9
         assert out["selector_metrics"]["selector_coverage"].item() > 0.9
 
+    @pytest.mark.parametrize("lk_loss_type", ["ce", "lambda"])
+    def test_fused_loss_matches_the_materialized_logits(self, lk_loss_type):
+        """The selector's top-k and its gradient come out of the fused pass unchanged."""
+        results = []
+        for fused in (False, True):
+            torch.manual_seed(0)
+            model = get_tiny_llama(num_hidden_layers=4)
+            config = _get_dflash2_config()
+            config.update(dflash_fused_loss=fused, dflash_lk_loss_type=lk_loss_type)
+            mtsp.convert(model, [("dflash", config)])
+            model.float().train()  # fp32, so the two paths differ only by rounding
+            with torch.no_grad():  # a live selector, so its gradient reaches the backbone
+                torch.nn.init.normal_(model.dflash_module.candidate_selector.successor_codebook)
+            input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size)
+            torch.manual_seed(1)
+            out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+            out.loss.backward()
+            params = model.dflash_module.named_parameters()
+            results.append((out, {name: p.grad for name, p in params if p.grad is not None}))
+        (eager, eager_grads), (fused, fused_grads) = results
+        torch.testing.assert_close(fused.loss, eager.loss)
+        assert fused.train_acc == eager.train_acc
+        for key, value in eager["selector_metrics"].items():
+            torch.testing.assert_close(fused["selector_metrics"][key], value, msg=key)
+        assert fused_grads.keys() == eager_grads.keys()
+        for name, grad in eager_grads.items():
+            torch.testing.assert_close(fused_grads[name], grad, msg=name)
+
 
 class TestCandidateSelectorAlignment:
     """Pin the offset convention shared by the training objective and the lattice walk."""
diff --git a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
index 0d132f9ce49..696b7b6ac36 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
@@ -428,6 +428,31 @@ def test_offline_forward_matches_online(self, prenorm):
                 online_out.lilicorr_metrics[term], abs=1e-5
             )
 
+    def test_fused_loss_matches_the_materialized_logits(self):
+        """The candidates, their logits and the normalizer come out of the fused pass unchanged."""
+        results = []
+        for fused in (False, True):
+            torch.manual_seed(0)
+            # Half the vocabulary per slot: the lattice covers about half the blocks.
+            model = _converted(
+                variant="margin", topk=16, dflash_lilicorr_w_cal=0.5, dflash_fused_loss=fused
+            )
+            model.float().train()  # fp32, so the two paths differ only by rounding
+            batch = _make_batch(model.dflash_config.vocab_size)
+            torch.manual_seed(1)
+            out = model(**batch)
+            out.loss.backward()
+            params = model.dflash_module.named_parameters()
+            results.append((out, {name: p.grad for name, p in params if p.grad is not None}))
+        (eager, eager_grads), (fused, fused_grads) = results
+        torch.testing.assert_close(fused.loss, eager.loss)
+        assert fused.lilicorr_metrics == pytest.approx(eager.lilicorr_metrics, rel=1e-5, abs=1e-6)
+        assert fused_grads.keys() == eager_grads.keys()
+        for name, grad in eager_grads.items():
+            # Scaled to each gradient: the normalizer's share of it is below a fixed atol.
+            atol = 1e-5 * grad.abs().max().item()
+            torch.testing.assert_close(fused_grads[name], grad, rtol=0, atol=atol, msg=name)
+
     @pytest.mark.parametrize(
         ("kwargs", "degenerate"),
         [