diff --git a/docs/guides/speculative/dflash.mdx b/docs/guides/speculative/dflash.mdx
index 6bfd159f66..e4b610f770 100644
--- a/docs/guides/speculative/dflash.mdx
+++ b/docs/guides/speculative/dflash.mdx
@@ -113,8 +113,11 @@ and `val_accept_len`; adding a top-level `wandb:` block uploads all of them.
 | `draft_num_hidden_layers` | draft stack depth (paper: 5) |
 | `block_size` | tokens drafted in parallel per block (paper: 16) |
 | `num_anchors` | blocks sampled per sequence per step |
+| `max_total_anchors` | caps sampled anchor blocks across the local micro-batch to bound memory use |
 | `loss_decay_gamma` | block-position decay; `null` disables it |
 | `loss_type` | `dflash` (fixed anchor) or `variable_prefix` |
+| `use_fused_linear_ce` | chunks the LM-head projection and cross-entropy instead of materializing all logits; supports `loss_type: dflash` only |
+| `linear_ce_chunk_size` | positions projected per fused CE chunk; smaller values reduce peak memory further |
 | `target_layer_ids` | target feature layers fed to the draft (defaults to an even spread) |
 | `mask_token_id` | reserved token id filling non-anchor block positions (required) |
 | `attention_backend` | `flex_attention` (GPU main path) or `sdpa` (portable fallback) |
diff --git a/examples/speculative/dflash/qwen3_dflash.yaml b/examples/speculative/dflash/qwen3_dflash.yaml
index 23c4cfdc86..075f561e33 100644
--- a/examples/speculative/dflash/qwen3_dflash.yaml
+++ b/examples/speculative/dflash/qwen3_dflash.yaml
@@ -29,8 +29,11 @@ recipe_args:
   # --- DFlash-specific ---
   draft_num_hidden_layers: 5          # depth of the draft transformer stack
   block_size: 16                      # tokens drafted in parallel per block
-  num_anchors: 512                    # blocks sampled per sequence per step
+  num_anchors: 512                    # per-sequence candidate block limit
+  max_total_anchors: 512              # cap constructed block slots across the local microbatch
   loss_decay_gamma: 7.0               # block-position loss decay: exp(-(k-1)/gamma)
+  use_fused_linear_ce: true            # chunk LM-head + CE; dflash loss only (disable for variable_prefix)
+  linear_ce_chunk_size: 256            # projected positions per chunk; lower uses less peak memory
   # loss_type: dflash                 # "dflash" (fixed anchor) or "variable_prefix" (D2SD VP-Drafter:
   #                                   # per-block visible prefix sampled ~ prefix_weight_base**l, only the
   #                                   # masked suffix is supervised, decay re-anchored at the prefix)
diff --git a/examples/speculative/dflash/qwen3_dflash_tp.yaml b/examples/speculative/dflash/qwen3_dflash_tp.yaml
index ea9ac9b988..fd3d5ef54a 100644
--- a/examples/speculative/dflash/qwen3_dflash_tp.yaml
+++ b/examples/speculative/dflash/qwen3_dflash_tp.yaml
@@ -36,8 +36,11 @@ recipe_args:
   # --- DFlash-specific ---
   draft_num_hidden_layers: 5
   block_size: 16
-  num_anchors: 512
+  num_anchors: 512                    # per-sequence candidate block limit
+  max_total_anchors: 512              # cap constructed block slots across the local microbatch
   loss_decay_gamma: 7.0
+  use_fused_linear_ce: true            # dflash loss only; disable when selecting variable_prefix
+  linear_ce_chunk_size: 256
   mask_token_id: 151669               # REQUIRED reserved token for non-anchor block positions
   attention_backend: flex_attention
 
diff --git a/nemo_automodel/components/loss/dllm_loss.py b/nemo_automodel/components/loss/dllm_loss.py
index 504aaabd94..d3f4e3462d 100644
--- a/nemo_automodel/components/loss/dllm_loss.py
+++ b/nemo_automodel/components/loss/dllm_loss.py
@@ -818,7 +818,24 @@ def _reduce(
         draft_correct_per_pos: torch.Tensor | None = None,
         draft_count_per_pos: torch.Tensor | None = None,
     ) -> DLLMLossOutput:
-        """Apply decay weights + block mask, sum, and normalise."""
+        """Apply decay weights and the block mask, then normalize the loss.
+
+        Args:
+            token_nll: Tensor of shape ``[batch, tokens]`` containing per-token
+                negative log-likelihood values.
+            block_mask: Tensor of shape ``[batch, tokens]`` containing valid-token
+                indicators.
+            num_tokens: Optional global token-count denominator.
+            block_size: Optional size of each draft block, including its anchor.
+            draft_correct_per_pos: Optional tensor of shape ``[block_size - 1]``
+                containing correct-prediction counts by draft depth.
+            draft_count_per_pos: Optional tensor of shape ``[block_size - 1]``
+                containing valid-token counts by draft depth.
+
+        Returns:
+            DLLMLossOutput containing scalar loss tensors and optional per-depth
+            count tensors.
+        """
         _, T = token_nll.shape
         w = self._decay_weights(T, block_size, token_nll.device, token_nll.dtype)
         weights = w.unsqueeze(0) * block_mask.to(token_nll.dtype)  # [B, T]
@@ -911,13 +928,23 @@ def _chunk_nll(
         Wrapped in :func:`torch.utils.checkpoint` by the caller, so the
         ``[chunk, vocab]`` logits are recomputed in backward rather than held.
         The argmax is non-differentiable, so it adds no backward cost.
+
+        Args:
+            hidden_chunk: Tensor of shape ``[positions, hidden]``.
+            lm_head_weight: Tensor of shape ``[vocab, hidden]``.
+            lm_head_bias: Optional tensor of shape ``[vocab]``.
+            target_chunk: Long tensor of shape ``[positions]``.
+
+        Returns:
+            Tuple containing per-token NLL and boolean argmax correctness tensors,
+            each of shape ``[positions]``.
         """
         logits = F.linear(hidden_chunk, lm_head_weight, lm_head_bias)  # [chunk, V]
         nll = F.cross_entropy(logits.float(), target_chunk, reduction="none")  # [chunk]
         correct = logits.argmax(dim=-1) == target_chunk  # [chunk]
         return nll, correct
 
-    def forward_fused(
+    def forward_fused_with_correct(
         self,
         hidden: torch.Tensor,
         lm_head_weight: torch.Tensor,
@@ -926,27 +953,30 @@ def forward_fused(
         num_tokens: int | None = None,
         block_size: int | None = None,
         lm_head_bias: torch.Tensor | None = None,
-    ) -> DLLMLossOutput:
-        """Chunked linear-CE: never materialises the full logits tensor.
+    ) -> tuple[DLLMLossOutput, torch.Tensor]:
+        """Compute chunked linear-CE and retain per-token correctness.
 
-        Projects the LM head + cross-entropy in chunks of ``chunk_size``
-        predicted positions, each wrapped in :func:`torch.utils.checkpoint` so
-        the ``[chunk, vocab]`` logits are recomputed in backward instead of
-        held — peak logit memory is one chunk, not ``[B*T, vocab]``. Pure
-        autograd, so the gradient flows correctly through FSDP2 (unlike a
-        standalone liger fused-CE Function).
+        This is the detailed variant used by DFlash trainers that need acceptance
+        metrics. :meth:`forward_fused` preserves the historical four-field
+        :class:`DLLMLossOutput` contract by discarding the detailed correctness
+        tensor.
 
         Args:
-            hidden: Draft hidden states for the predicted positions,
-                shape ``[B, T, D]`` (``D`` = model dim, NOT vocab).
-            lm_head_weight: LM-head projection weight, shape ``[V, D]``.
-            target_ids: Ground-truth token IDs, shape ``[B, T]``.
-            block_mask: Valid-position mask, shape ``[B, T]``.
-            num_tokens / block_size: as in :meth:`forward`.
-            lm_head_bias: Optional LM-head bias, shape ``[V]``.
+            hidden: Tensor of shape ``[batch, tokens, hidden]`` containing draft
+                hidden states.
+            lm_head_weight: Plain local tensor of shape ``[vocab, hidden]``. Any
+                distributed weight must be materialized by the caller before this
+                method so the position-chunk loop contains no collectives.
+            target_ids: Long tensor of shape ``[batch, tokens]``.
+            block_mask: Tensor of shape ``[batch, tokens]`` containing valid-token
+                indicators.
+            num_tokens: Optional global token-count denominator.
+            block_size: Optional size of each draft block, including its anchor.
+            lm_head_bias: Optional plain local tensor of shape ``[vocab]``.
 
         Returns:
-            :class:`DLLMLossOutput`.
+            Tuple containing a DLLMLossOutput and a boolean correctness tensor of
+            shape ``[batch, tokens]``.
         """
         B, T, D = hidden.shape
         flat_hidden = hidden.reshape(-1, D)  # [B*T, D]
@@ -969,7 +999,7 @@ def forward_fused(
         token_nll = torch.cat(nll_parts).reshape(B, T)
         correct = torch.cat(correct_parts).reshape(B, T)
         c_per_pos, n_per_pos = self._draft_acc_per_pos(correct, block_mask, block_size)
-        return self._reduce(
+        output = self._reduce(
             token_nll,
             block_mask,
             num_tokens,
@@ -977,6 +1007,49 @@ def forward_fused(
             draft_correct_per_pos=c_per_pos,
             draft_count_per_pos=n_per_pos,
         )
+        return output, correct
+
+    def forward_fused(
+        self,
+        hidden: torch.Tensor,
+        lm_head_weight: torch.Tensor,
+        target_ids: torch.Tensor,
+        block_mask: torch.Tensor,
+        num_tokens: int | None = None,
+        block_size: int | None = None,
+        lm_head_bias: torch.Tensor | None = None,
+    ) -> DLLMLossOutput:
+        """Chunked linear-CE: never materialises the full logits tensor.
+
+        Projects the LM head + cross-entropy in chunks of ``chunk_size``
+        predicted positions, each wrapped in :func:`torch.utils.checkpoint` so
+        the ``[chunk, vocab]`` logits are recomputed in backward instead of
+        held — peak logit memory is one chunk, not ``[B*T, vocab]``. Pure
+        autograd, so the gradient flows correctly through FSDP2 (unlike a
+        standalone liger fused-CE Function).
+
+        Args:
+            hidden: Draft hidden states for the predicted positions,
+                shape ``[B, T, D]`` (``D`` = model dim, NOT vocab).
+            lm_head_weight: LM-head projection weight, shape ``[V, D]``.
+            target_ids: Ground-truth token IDs, shape ``[B, T]``.
+            block_mask: Valid-position mask, shape ``[B, T]``.
+            num_tokens / block_size: as in :meth:`forward`.
+            lm_head_bias: Optional LM-head bias, shape ``[V]``.
+
+        Returns:
+            :class:`DLLMLossOutput`.
+        """
+        output, _ = self.forward_fused_with_correct(
+            hidden,
+            lm_head_weight,
+            target_ids,
+            block_mask,
+            num_tokens=num_tokens,
+            block_size=block_size,
+            lm_head_bias=lm_head_bias,
+        )
+        return output
 
 
 class IDLMLoss(nn.Module):
diff --git a/nemo_automodel/components/speculative/dflash/core.py b/nemo_automodel/components/speculative/dflash/core.py
index d2d015a546..fb2bf897bf 100644
--- a/nemo_automodel/components/speculative/dflash/core.py
+++ b/nemo_automodel/components/speculative/dflash/core.py
@@ -68,12 +68,15 @@ def _context_doc_ids(seq_lens: torch.Tensor, seq_len: int, device: torch.device)
 
 
 def _to_full_tensor(tensor: torch.Tensor) -> torch.Tensor:
-    """Materialise a (possibly tensor-parallel) tensor as a plain local tensor.
+    """Materialize a distributed tensor as a plain local tensor.
 
-    Under tensor parallelism the target's column-parallel ``lm_head`` and
-    vocab-parallel ``embed_tokens`` return ``DTensor`` outputs. The draft and the
-    block-wise loss consume plain tensors, so gather the full tensor. A no-op for
-    an already-plain (unsharded / replicated) tensor.
+    Args:
+        tensor: Tensor of arbitrary shape. A DTensor may be sharded over one or
+            more mesh axes; a plain tensor is returned unchanged.
+
+    Returns:
+        Plain local tensor with the DTensor's global shape, or the original plain
+        tensor without copying.
     """
     return tensor.full_tensor() if hasattr(tensor, "full_tensor") else tensor
 
@@ -157,7 +160,26 @@ def compute_acceptance_stats(
         Three scalar tensors: mean acceptance length, its additive sum, and the
         number of blocks containing at least one valid drafted token.
     """
-    block_accept = compute_accept_len(pred_ids_4d, target_ids_4d, valid_mask_4d)
+    correct_4d = pred_ids_4d == target_ids_4d
+    return _compute_acceptance_stats_from_correct(correct_4d, valid_mask_4d)
+
+
+def _compute_acceptance_stats_from_correct(
+    correct_4d: torch.Tensor,
+    valid_mask_4d: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Return acceptance statistics from per-token correctness.
+
+    Args:
+        correct_4d: Bool tensor of shape ``[batch, blocks, depth]``.
+        valid_mask_4d: Bool tensor of shape ``[batch, blocks, depth]``.
+
+    Returns:
+        Three scalar tensors containing mean acceptance length, additive
+        acceptance-length sum, and valid-block count.
+    """
+    correct_or_invalid = correct_4d | (~valid_mask_4d)
+    block_accept = (correct_or_invalid.long().cumprod(dim=2) * valid_mask_4d.long()).sum(dim=2).float()
     valid_block_mask = valid_mask_4d.any(dim=2)
     valid_blocks = valid_block_mask.sum()
     accept_len_sum = ((block_accept + 1.0) * valid_block_mask).sum()
@@ -169,7 +191,25 @@ def compute_acceptance_stats(
 
 
 class DFlashTrainerModule(nn.Module):
-    """DFlash online training wrapper with block-wise CE loss."""
+    """DFlash online training wrapper with block-wise CE loss.
+
+    Args:
+        draft_model: Trainable DFlash draft model.
+        target_lm_head: Frozen target token-projection module.
+        target_embed_tokens: Frozen target token-embedding module.
+        mask_token_id: Token ID used for masked draft positions.
+        block_size: Number of tokens in each draft block, including its anchor.
+        attention_backend: Draft attention-mask backend.
+        num_anchors: Per-sequence candidate-anchor limit.
+        loss_decay_gamma: Optional exponential depth-decay parameter.
+        loss_type: DFlash training objective name.
+        prefix_weight_base: Truncated-geometric base for variable-prefix training.
+        max_total_anchors: Optional cap on rectangular anchor slots across the
+            local microbatch.
+        use_fused_linear_ce: Whether to materialize the frozen target projection
+            once and run chunked linear cross-entropy.
+        linear_ce_chunk_size: Number of predicted positions projected per chunk.
+    """
 
     def __init__(
         self,
@@ -184,12 +224,24 @@ def __init__(
         loss_type: str = "dflash",
         prefix_weight_base: float = 0.9,
         sliding_window: int | None = None,
+        *,
+        max_total_anchors: int | None = None,
+        use_fused_linear_ce: bool = False,
+        linear_ce_chunk_size: int = 1024,
     ):
         super().__init__()
         if loss_type not in _DFLASH_LOSS_TYPES:
             raise ValueError(f"loss_type must be one of {_DFLASH_LOSS_TYPES}, got {loss_type!r}")
         if prefix_weight_base <= 0:
             raise ValueError(f"prefix_weight_base must be > 0, got {prefix_weight_base}")
+        if num_anchors <= 0:
+            raise ValueError(f"num_anchors must be > 0, got {num_anchors}")
+        if max_total_anchors is not None and max_total_anchors <= 0:
+            raise ValueError(f"max_total_anchors must be > 0 or None, got {max_total_anchors}")
+        if linear_ce_chunk_size <= 0:
+            raise ValueError(f"linear_ce_chunk_size must be > 0, got {linear_ce_chunk_size}")
+        if use_fused_linear_ce and loss_type != "dflash":
+            raise ValueError("use_fused_linear_ce is only supported with loss_type='dflash'")
         self.draft_model = draft_model
         # Keep the frozen target lm_head / embed_tokens as NON-registered
         # references. Under tensor parallelism their weights are DTensors; a
@@ -210,9 +262,11 @@ def __init__(
         if sliding_window is not None and sliding_window < 1:
             raise ValueError(f"sliding_window must be >= 1 when set, got {sliding_window}.")
         self.sliding_window = sliding_window
+        self.max_total_anchors = max_total_anchors
         self.loss_decay_gamma = loss_decay_gamma
         self.loss_type = loss_type
         self.prefix_weight_base = float(prefix_weight_base)
+        self.use_fused_linear_ce = bool(use_fused_linear_ce)
         # Smallest visible-prefix length variable-prefix training samples (and the
         # slice point of its loss); single source of truth for both methods.
         self._min_prefix = min(2, block_size - 1)
@@ -222,11 +276,61 @@ def __init__(
         # ``loss_decay_gamma=None`` disables decay (uniform weights). The
         # variable-prefix objective needs per-block data-dependent weights and
         # computes its loss inline instead (see _variable_prefix_loss).
-        self.loss_fn = DFlashDecayLoss(loss_gamma=loss_decay_gamma, normalize="mean") if loss_type == "dflash" else None
+        self.loss_fn = (
+            DFlashDecayLoss(
+                loss_gamma=loss_decay_gamma,
+                use_fused_linear_ce=use_fused_linear_ce,
+                chunk_size=linear_ce_chunk_size,
+                normalize="mean",
+            )
+            if loss_type == "dflash"
+            else None
+        )
 
         # Per-block offset constant (block_size,) for label gathering / position ids.
         self.register_buffer("_block_offsets", torch.arange(block_size).view(1, 1, -1), persistent=False)
 
+    def _materialize_frozen_lm_head(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Gather the frozen target projection once before chunked linear-CE.
+
+        The target head may be a child whose parameter is owned by an ancestor
+        FSDP2 unit, so invoking the child module cannot reliably trigger the owning
+        unit's unshard hook. Materializing its DTensor parameters directly is an
+        explicit rank-symmetric collective contract: every participating rank
+        gathers once before anchor sampling, then every position chunk uses plain
+        local tensors and performs no distributed collective.
+
+        Args:
+            device: Device where fused linear-CE will consume the projection.
+                This may differ from the frozen FSDP shard's storage device when
+                CPU offload is enabled or the head is not exercised by the
+                target hidden-state forward.
+
+        Returns:
+            Tuple containing a plain tensor of shape ``[vocab, hidden]`` and an
+            optional plain bias tensor of shape ``[vocab]``, both on ``device``.
+
+        Raises:
+            TypeError: If the target head does not expose tensor weight/bias fields.
+            ValueError: If the target projection is trainable rather than frozen.
+        """
+        weight = getattr(self.lm_head, "weight", None)
+        bias = getattr(self.lm_head, "bias", None)
+        if not isinstance(weight, torch.Tensor):
+            raise TypeError("Fused DFlash linear-CE requires target_lm_head.weight to be a tensor")
+        if bias is not None and not isinstance(bias, torch.Tensor):
+            raise TypeError("Fused DFlash linear-CE requires target_lm_head.bias to be a tensor or None")
+        # Root-owned FSDP2 can expose a zero-sized placeholder for a bias-free
+        # child projection. Passing that placeholder to F.linear is not
+        # equivalent to ``bias=None`` and fails during broadcast.
+        if bias is not None and bias.numel() == 0:
+            bias = None
+        if weight.requires_grad or (bias is not None and bias.requires_grad):
+            raise ValueError("Fused DFlash linear-CE requires a frozen target LM head")
+        full_weight = _to_full_tensor(weight).to(device=device)
+        full_bias = _to_full_tensor(bias).to(device=device) if bias is not None else None
+        return full_weight, full_bias
+
     def _sample_anchor_positions(
         self,
         seq_len: int,
@@ -263,6 +367,15 @@ def _sample_anchor_positions(
         # by ``keep_mask`` below); no -1, which would spuriously raise when the
         # richest sample has exactly one valid anchor and always drop one otherwise.
         max_n = min(self.num_anchors, int(valid_counts.max().item()))
+        if self.max_total_anchors is not None:
+            if self.max_total_anchors < bsz:
+                raise ValueError(
+                    f"max_total_anchors ({self.max_total_anchors}) must be at least the local batch size ({bsz})"
+                )
+            # Blocks are represented by rectangular [B, N, ...] tensors. Cap N
+            # uniformly so the number of constructed slots B*N stays within the
+            # local micro-batch budget, including padded slots.
+            max_n = min(max_n, self.max_total_anchors // bsz)
         if max_n <= 0:
             doc_note = " with block_size-1 further real tokens in its document" if doc_remaining is not None else ""
             raise NoValidAnchorsError(
@@ -531,9 +644,34 @@ def forward(
         keeps every block inside one document: anchors are constrained so the block
         does not cross a boundary, the block's context prefix attends only within the
         anchor's document, and the draft's RoPE uses the per-document positions.
+
+        Args:
+            input_ids: Long tensor of shape ``[batch, sequence]``.
+            hidden_states: Target-conditioning tensor of shape ``[batch, sequence,
+                conditioning_hidden]``.
+            loss_mask: Tensor of shape ``[batch, sequence]`` containing supervised-
+                token indicators.
+            position_ids: Optional long tensor of shape ``[batch, sequence]`` with
+                per-document positions for packed input.
+            seq_lens: Optional long tensor of shape ``[batch, documents]`` containing
+                packed-document lengths.
+            doc_remaining: Optional long tensor of shape ``[batch, sequence]``
+                containing remaining tokens in each position's document.
+
+        Returns:
+            DFlashStepMetrics containing scalar loss, accuracy, token-count, and
+            acceptance-length tensors.
         """
         bsz, seq_len = input_ids.shape
 
+        # Materialize before data-dependent anchor validation. With an FSDP-owned
+        # head this may enter collectives, so every rank must do so in the same
+        # order even when a later NoValidAnchorsError makes the recipe skip the
+        # microbatch. The returned full weight is reused by every CE chunk.
+        lm_head_weight, lm_head_bias = (
+            self._materialize_frozen_lm_head(input_ids.device) if self.use_fused_linear_ce else (None, None)
+        )
+
         anchor_positions, block_keep_mask, noise_embedding, full_position_ids, dflash_attn_mask, prefix_lengths = (
             self._prepare_block_inputs(
                 input_ids, loss_mask, position_ids=position_ids, seq_lens=seq_lens, doc_remaining=doc_remaining
@@ -546,10 +684,6 @@ def forward(
             target_hidden=hidden_states,
             attention_mask=dflash_attn_mask,
         )
-        # A tensor-parallel target's lm_head is column-parallel and returns
-        # vocab-sharded (DTensor) logits; gather to a full tensor for the loss.
-        logits = _to_full_tensor(self.lm_head(output_hidden))
-
         n = anchor_positions.size(1)
         bs = self.block_size
 
@@ -559,17 +693,36 @@ def forward(
         )
 
         if self.loss_type == "variable_prefix":
+            # Variable-prefix training has data-dependent per-block suffixes and
+            # still uses the dense objective.
+            logits = _to_full_tensor(self.lm_head(output_hidden))
             return self._variable_prefix_loss(logits.view(bsz, n, bs, -1), target_ids, block_mask, prefix_lengths)
 
         # Drop block position 0 (the clean anchor token, never a target); the
         # remaining bs-1 predicted positions are what the loss supervises.
-        pred_logits = logits.view(bsz, n, bs, -1)[:, :, 1:, :].reshape(bsz, n * (bs - 1), -1)
+        pred_hidden = output_hidden.view(bsz, n, bs, -1)[:, :, 1:, :].reshape(bsz, n * (bs - 1), -1)
         pred_targets = target_ids[:, :, 1:].reshape(bsz, n * (bs - 1))
         pred_mask = block_mask[:, :, 1:].reshape(bsz, n * (bs - 1))
 
         loss_fn = self.loss_fn
         assert loss_fn is not None, "loss_fn is always constructed for loss_type='dflash'"
-        loss_out = loss_fn(pred_logits, pred_targets, pred_mask, num_tokens=None, block_size=bs)
+        if self.use_fused_linear_ce:
+            assert lm_head_weight is not None, "the fused path always materializes the target LM-head weight"
+            loss_out, draft_correct = loss_fn.forward_fused_with_correct(
+                hidden=pred_hidden,
+                lm_head_weight=lm_head_weight,
+                target_ids=pred_targets,
+                block_mask=pred_mask,
+                num_tokens=None,
+                block_size=bs,
+                lm_head_bias=lm_head_bias,
+            )
+        else:
+            # A tensor-parallel target's lm_head returns vocab-sharded DTensor
+            # logits; gather the full tensor for the dense fallback loss.
+            pred_logits = _to_full_tensor(self.lm_head(pred_hidden))
+            loss_out = loss_fn(pred_logits, pred_targets, pred_mask, num_tokens=None, block_size=bs)
+            draft_correct = pred_logits.argmax(dim=-1) == pred_targets
 
         loss_weights = pred_mask.view(bsz, n, bs - 1)
         if self.loss_decay_gamma is not None:
@@ -580,12 +733,13 @@ def forward(
         loss_weight = loss_weights.sum()
 
         count_per_pos = loss_out.draft_count_per_pos
+        correct_per_pos = loss_out.draft_correct_per_pos
+        assert count_per_pos is not None and correct_per_pos is not None
         valid_tokens = count_per_pos.sum()
-        correct_tokens = loss_out.draft_correct_per_pos.sum()
+        correct_tokens = correct_per_pos.sum()
         accuracy = correct_tokens / valid_tokens.clamp_min(1)
-        accept_len, accept_len_sum, valid_blocks = compute_acceptance_stats(
-            pred_logits.argmax(dim=-1).view(bsz, n, bs - 1),
-            pred_targets.view(bsz, n, bs - 1),
+        accept_len, accept_len_sum, valid_blocks = _compute_acceptance_stats_from_correct(
+            draft_correct.view(bsz, n, bs - 1),
             pred_mask.view(bsz, n, bs - 1).bool(),
         )
 
diff --git a/nemo_automodel/components/speculative/dflash/dflash2_core.py b/nemo_automodel/components/speculative/dflash/dflash2_core.py
index 92ee00b837..6de50aafc1 100644
--- a/nemo_automodel/components/speculative/dflash/dflash2_core.py
+++ b/nemo_automodel/components/speculative/dflash/dflash2_core.py
@@ -75,6 +75,8 @@ class DFlash2StepMetrics:
         valid_blocks: Scalar tensor containing the number of evaluated draft blocks.
         base_loss: Scalar tensor containing the backbone block-CE term.
         selector_loss: Scalar tensor containing the candidate-selection CE term.
+        selector_loss_denominator: Scalar tensor containing the effective
+            denominator for the candidate-selection CE term.
         base_accuracy: Scalar tensor containing backbone top-1 token accuracy.
         base_correct_tokens: Scalar tensor containing backbone correct-token count.
         base_accept_len: Scalar tensor containing backbone mean acceptance length.
@@ -94,6 +96,7 @@ class DFlash2StepMetrics:
     valid_blocks: torch.Tensor
     base_loss: torch.Tensor
     selector_loss: torch.Tensor
+    selector_loss_denominator: torch.Tensor
     base_accuracy: torch.Tensor
     base_correct_tokens: torch.Tensor
     base_accept_len: torch.Tensor
@@ -116,6 +119,8 @@ def __init__(
         loss_decay_gamma: float | None = None,
         selector_loss_weight: float = 1.0,
         sliding_window: int | None = None,
+        *,
+        max_total_anchors: int | None = None,
     ):
         super().__init__(
             draft_model=draft_model,
@@ -127,6 +132,7 @@ def __init__(
             num_anchors=num_anchors,
             loss_decay_gamma=loss_decay_gamma,
             sliding_window=sliding_window,
+            max_total_anchors=max_total_anchors,
         )
         if getattr(draft_model, "candidate_selector", None) is None:
             raise ValueError(
@@ -290,8 +296,9 @@ def forward(
             accept_len=accept_len.detach(),
             accept_len_sum=accept_len_sum.detach(),
             valid_blocks=valid_blocks.detach(),
-            base_loss=loss_out.total_loss.detach(),
-            selector_loss=selector_loss.detach(),
+            base_loss=loss_out.total_loss,
+            selector_loss=selector_loss,
+            selector_loss_denominator=selector_weights.sum().detach(),
             base_accuracy=(base_correct_tokens / denominator).detach(),
             base_correct_tokens=base_correct_tokens.detach(),
             base_accept_len=base_accept_len.detach(),
diff --git a/nemo_automodel/components/speculative/dflash/domino_core.py b/nemo_automodel/components/speculative/dflash/domino_core.py
index 9b7302c891..7b7e9acbc0 100644
--- a/nemo_automodel/components/speculative/dflash/domino_core.py
+++ b/nemo_automodel/components/speculative/dflash/domino_core.py
@@ -131,6 +131,8 @@ def __init__(
         num_anchors: int = 512,
         loss_decay_gamma: float | None = None,
         shift_label: bool = False,
+        *,
+        max_total_anchors: int | None = None,
     ):
         super().__init__(
             draft_model=draft_model,
@@ -141,6 +143,7 @@ def __init__(
             attention_backend=attention_backend,
             num_anchors=num_anchors,
             loss_decay_gamma=loss_decay_gamma,
+            max_total_anchors=max_total_anchors,
         )
         if getattr(draft_model, "projector_type", None) != "domino":
             raise ValueError(
diff --git a/nemo_automodel/components/speculative/dflash/draft_qwen3.py b/nemo_automodel/components/speculative/dflash/draft_qwen3.py
index 3fe0d7c35d..e1a05c7267 100644
--- a/nemo_automodel/components/speculative/dflash/draft_qwen3.py
+++ b/nemo_automodel/components/speculative/dflash/draft_qwen3.py
@@ -29,7 +29,7 @@
 
 from __future__ import annotations
 
-from typing import Callable, Tuple
+from typing import Any, Callable, Tuple
 
 import torch
 from torch import nn
@@ -47,7 +47,7 @@
     rotate_half,
 )
 
-from nemo_automodel.components.speculative.dflash.target import resolve_text_config
+from nemo_automodel.components.speculative.dflash.target import resolve_text_config, resolve_transformer_layers
 
 
 def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor:
@@ -222,6 +222,15 @@ def forward(
         attn_fn: Callable = eager_attention_forward
         if self.config._attn_implementation != "eager":
             attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        if self.config._attn_implementation == "flex_attention":
+            # PyTorch's AUTO backend selects a separate flex-decoding lowering
+            # for rank-local query lengths below 128. Some dynamic speech
+            # batches have no valid decoding autotune choice, while other ranks
+            # select the regular FlexAttention kernel and proceed. Force the
+            # regular kernel so every rank supports the same dynamic shapes.
+            kernel_options = dict(kwargs.get("kernel_options") or {})
+            kernel_options.setdefault("FORCE_USE_FLEX_ATTENTION", True)
+            kwargs["kernel_options"] = kernel_options
         attn_output, attn_weights = attn_fn(
             self,
             q,
@@ -339,12 +348,73 @@ def extract_context_feature(hidden_states: list[torch.Tensor], layer_ids: list[i
 
     ``hidden_states`` follows HF's ``output_hidden_states`` convention where
     index 0 is the embedding output, so layer ``i``'s output is at index
-    ``i + 1``.
+    ``i + 1``. The final tuple entry may already include the model-level final
+    norm, so selecting the final decoder block through this compatibility helper
+    is rejected; generation uses forward hooks instead.
     """
+    final_block_id = len(hidden_states) - 2
+    if final_block_id in layer_ids:
+        raise ValueError(
+            "The final decoder block cannot be recovered pre-norm from output_hidden_states; "
+            "capture it with forward_target_with_context_feature instead."
+        )
     offset = 1
     return torch.cat([hidden_states[layer_id + offset] for layer_id in layer_ids], dim=-1)
 
 
+def forward_target_with_context_feature(
+    target: nn.Module,
+    input_ids: torch.LongTensor,
+    layer_ids: list[int],
+    **forward_kwargs: Any,
+) -> tuple[Any, torch.Tensor]:
+    """Run a target forward and capture raw decoder-block outputs.
+
+    Hugging Face replaces the final decoder block's entry in
+    ``output_hidden_states`` with the model-level normalized state. Forward
+    hooks preserve one feature contract for every configured block, including
+    the final one: its raw residual output before that separate normalization.
+
+    Args:
+        target: Frozen target language model.
+        input_ids: Target token IDs with shape ``[batch, sequence]``.
+        layer_ids: Decoder-block IDs to concatenate in the requested order.
+        **forward_kwargs: Additional keyword arguments for the target forward.
+
+    Returns:
+        Pair of the target model output and concatenated pre-final-norm context
+        features with shape ``[batch, sequence, selected_layers * hidden]``.
+
+    Raises:
+        ValueError: If a configured decoder-block ID is out of bounds.
+        RuntimeError: If a configured decoder block does not run.
+    """
+    layers = resolve_transformer_layers(target)
+    for layer_id in layer_ids:
+        if layer_id < 0 or layer_id >= len(layers):
+            raise ValueError(f"target layer id {layer_id} is out of bounds for model with {len(layers)} layers")
+    captured: dict[int, torch.Tensor] = {}
+    handles = []
+
+    def make_hook(layer_id: int):
+        def hook(_module, _inputs, outputs):
+            captured[layer_id] = outputs[0] if isinstance(outputs, tuple) else outputs
+
+        return hook
+
+    try:
+        for layer_id in layer_ids:
+            handles.append(layers[layer_id].register_forward_hook(make_hook(layer_id)))
+        output = target(input_ids, **forward_kwargs)
+    finally:
+        for handle in handles:
+            handle.remove()
+
+    if len(captured) != len(layer_ids):
+        raise RuntimeError(f"Expected {len(layer_ids)} captured layers but got {len(captured)}: {sorted(captured)}")
+    return output, torch.cat([captured[layer_id] for layer_id in layer_ids], dim=-1)
+
+
 class Qwen3DFlashDraftModel(Qwen3PreTrainedModel):
     """DFlash draft model: a small non-causal Qwen3 stack over ``[context | noise]``."""
 
@@ -489,17 +559,17 @@ def spec_generate(
         past_key_values_draft = DynamicCache(config=self.config)
 
         # Prefill the target on the prompt.
-        output = target(
+        output, target_hidden = forward_target_with_context_feature(
+            target,
             input_ids,
+            self.target_layer_ids,
             position_ids=position_ids[:, :num_input_tokens],
             past_key_values=past_key_values_target,
             use_cache=True,
             logits_to_keep=1,
-            output_hidden_states=True,
         )
         output_ids[:, :num_input_tokens] = input_ids
         output_ids[:, num_input_tokens : num_input_tokens + 1] = sample(output.logits, temperature)
-        target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids)
 
         start = num_input_tokens
         while start < max_length:
@@ -518,12 +588,13 @@ def spec_generate(
             past_key_values_draft.crop(start)
             block_output_ids[:, 1:] = sample(draft_logits)
 
-            output = target(
+            output, target_hidden = forward_target_with_context_feature(
+                target,
                 block_output_ids,
+                self.target_layer_ids,
                 position_ids=block_position_ids,
                 past_key_values=past_key_values_target,
                 use_cache=True,
-                output_hidden_states=True,
             )
             posterior = sample(output.logits, temperature)
             acceptance_length = (block_output_ids[:, 1:] == posterior[:, :-1]).cumprod(dim=1).sum(dim=1)[0].item()
@@ -531,9 +602,7 @@ def spec_generate(
             output_ids[:, start + acceptance_length + 1] = posterior[:, acceptance_length]
             start += acceptance_length + 1
             past_key_values_target.crop(start)
-            target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids)[
-                :, : acceptance_length + 1, :
-            ]
+            target_hidden = target_hidden[:, : acceptance_length + 1, :]
             if stop_token_ids is not None and any(
                 stop_id in output_ids[:, num_input_tokens:] for stop_id in stop_token_ids
             ):
diff --git a/nemo_automodel/components/speculative/dflash/draft_qwen3_dflash2.py b/nemo_automodel/components/speculative/dflash/draft_qwen3_dflash2.py
index 5bde5f46a8..c7ad4a4c5c 100644
--- a/nemo_automodel/components/speculative/dflash/draft_qwen3_dflash2.py
+++ b/nemo_automodel/components/speculative/dflash/draft_qwen3_dflash2.py
@@ -64,7 +64,7 @@
     Qwen3DFlashDecoderLayer,
     Qwen3DFlashDraftModel,
     assert_target_supports_rollback,
-    extract_context_feature,
+    forward_target_with_context_feature,
     sample,
 )
 
@@ -551,7 +551,8 @@ def spec_generate(
 
         Args:
             target: The frozen verifier; must expose ``model.embed_tokens``,
-                ``lm_head``, and an HF-style forward with ``output_hidden_states``.
+                ``lm_head``, a resolvable decoder-block stack, and an HF-style
+                forward returning logits.
             input_ids: Long tensor of shape [1, prompt].
             max_new_tokens: Maximum number of tokens to generate.
             stop_token_ids: Token ids that end generation, or ``None``.
@@ -581,17 +582,17 @@ def spec_generate(
         past_key_values_draft = DynamicCache(config=self.config)
 
         # Prefill the target on the prompt.
-        output = target(
+        output, target_hidden = forward_target_with_context_feature(
+            target,
             input_ids,
+            self.target_layer_ids,
             position_ids=position_ids[:, :num_input_tokens],
             past_key_values=past_key_values_target,
             use_cache=True,
             logits_to_keep=1,
-            output_hidden_states=True,
         )
         output_ids[:, :num_input_tokens] = input_ids
         output_ids[:, num_input_tokens : num_input_tokens + 1] = sample(output.logits, temperature)
-        target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids)
 
         start = num_input_tokens
         while start < max_length:
@@ -614,12 +615,13 @@ def spec_generate(
             )
             block_output_ids[:, 1:] = draft_tokens
 
-            output = target(
+            output, target_hidden = forward_target_with_context_feature(
+                target,
                 block_output_ids,
+                self.target_layer_ids,
                 position_ids=block_position_ids,
                 past_key_values=past_key_values_target,
                 use_cache=True,
-                output_hidden_states=True,
             )
             if temperature > 0:
                 target_probs = torch.softmax(output.logits.float() / temperature, dim=-1)
@@ -634,9 +636,7 @@ def spec_generate(
             output_ids[:, start + acceptance_length + 1] = bonus
             start += acceptance_length + 1
             past_key_values_target.crop(start)
-            target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids)[
-                :, : acceptance_length + 1, :
-            ]
+            target_hidden = target_hidden[:, : acceptance_length + 1, :]
             if stop_token_ids is not None and any(
                 stop_id in output_ids[:, num_input_tokens:] for stop_id in stop_token_ids
             ):
diff --git a/nemo_automodel/components/speculative/dflash/jetspec_core.py b/nemo_automodel/components/speculative/dflash/jetspec_core.py
index 209866c9da..654124ec67 100644
--- a/nemo_automodel/components/speculative/dflash/jetspec_core.py
+++ b/nemo_automodel/components/speculative/dflash/jetspec_core.py
@@ -105,6 +105,8 @@ def __init__(
         num_anchors: int = 512,
         kd_temperature: float = 1.0,
         kd_chunk_size: int = 0,
+        *,
+        max_total_anchors: int | None = None,
     ):
         super().__init__(
             draft_model=draft_model,
@@ -115,6 +117,7 @@ def __init__(
             attention_backend=attention_backend,
             num_anchors=num_anchors,
             loss_decay_gamma=None,
+            max_total_anchors=max_total_anchors,
         )
         # Forward KL(P_target || Q_draft) with Hinton T^2 scaling and uniform
         # weighting over active draft positions (paper Eq. 9, no depth weighting in
diff --git a/nemo_automodel/components/speculative/dflash/target.py b/nemo_automodel/components/speculative/dflash/target.py
index fc614c9925..1c0024385f 100644
--- a/nemo_automodel/components/speculative/dflash/target.py
+++ b/nemo_automodel/components/speculative/dflash/target.py
@@ -71,12 +71,37 @@ def resolve_text_config(config: Any) -> Any:
     return getattr(config, "text_config", None) or config
 
 
+def resolve_transformer_layers(model: nn.Module) -> list[nn.Module]:
+    """Return decoder blocks as an ordered, integer-indexable list.
+
+    Args:
+        model: Target language model containing a supported decoder stack.
+
+    Returns:
+        Decoder blocks in forward order.
+
+    Raises:
+        ValueError: If the decoder stack cannot be located.
+    """
+    if hasattr(model, "model") and hasattr(model.model, "layers"):
+        container = model.model.layers
+    elif hasattr(model, "layers"):
+        container = model.layers
+    elif hasattr(model, "transformer") and hasattr(model.transformer, "h"):
+        container = model.transformer.h
+    else:
+        raise ValueError("Unsupported model structure for DFlash hidden-state capture")
+    if isinstance(container, nn.ModuleDict):
+        return [container[str(i)] for i in range(len(container))]
+    return list(container)
+
+
 class HFDFlashTargetModel:
     """Capture a set of decoder-layer hidden states from a frozen HF causal LM.
 
-    A forward hook on decoder layer ``i`` captures that layer's output, which in
-    HuggingFace's ``output_hidden_states`` convention is ``hidden_states[i + 1]``
-    -- matching SpecForge's ``extract_context_feature`` (offset 1).
+    A forward hook on decoder block ``i`` captures that block's raw residual
+    output. This is the pre-final-norm feature: a separate model-level final
+    normalization, when present, is intentionally excluded.
     """
 
     def __init__(
@@ -126,17 +151,7 @@ def _validate_layer_ids(self, target_layer_ids: Sequence[int]) -> list[int]:
 
     def _get_transformer_layers(self) -> list[nn.Module]:
         """Return decoder layers as an ordered, integer-indexable list."""
-        if hasattr(self.model, "model") and hasattr(self.model.model, "layers"):
-            container = self.model.model.layers
-        elif hasattr(self.model, "layers"):
-            container = self.model.layers
-        elif hasattr(self.model, "transformer") and hasattr(self.model.transformer, "h"):
-            container = self.model.transformer.h
-        else:
-            raise ValueError("Unsupported model structure for DFlash hidden-state capture")
-        if isinstance(container, nn.ModuleDict):
-            return [container[str(i)] for i in range(len(container))]
-        return list(container)
+        return resolve_transformer_layers(self.model)
 
     def get_input_embeddings(self) -> nn.Embedding:
         """Return the target model input embeddings."""
diff --git a/nemo_automodel/recipes/llm/train_dflash.py b/nemo_automodel/recipes/llm/train_dflash.py
index fdd24ccb36..abc7940805 100644
--- a/nemo_automodel/recipes/llm/train_dflash.py
+++ b/nemo_automodel/recipes/llm/train_dflash.py
@@ -144,7 +144,7 @@ def _validate_packing_gates(*, cp_size: int, target_attn_impl: str, micro_batch_
         )
 
 
-def _all_ranks_have_valid(local_has_valid: int, is_ddp: bool, device) -> bool:
+def _all_ranks_have_valid(local_has_valid: int, is_ddp: bool, device, process_group=None) -> bool:
     """Min-reduce a per-rank "this micro-batch has valid anchors" flag.
 
     Under DDP a data-dependent ``NoValidAnchorsError`` skip is per-rank: if one
@@ -158,10 +158,47 @@ def _all_ranks_have_valid(local_has_valid: int, is_ddp: bool, device) -> bool:
     if not is_ddp or not (dist.is_available() and dist.is_initialized()):
         return bool(local_has_valid)
     flag = torch.tensor([local_has_valid], device=device, dtype=torch.int32)
-    dist.all_reduce(flag, op=dist.ReduceOp.MIN)
+    dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=process_group)
     return bool(flag.item())
 
 
+def _normalize_ddp_loss(
+    loss: torch.Tensor,
+    loss_weight: torch.Tensor,
+    is_ddp: bool,
+    process_group=None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Weight a local mean so DDP produces the global weighted-mean gradient.
+
+    DFlash micro-batches can contain different numbers of valid weighted anchor
+    positions on different data-parallel ranks. DDP averages rank gradients, so
+    backpropagating each rank's local mean directly would give every rank equal
+    weight. Scaling by ``local_weight * dp_size / global_weight`` compensates for
+    that averaging. The second return value is the detached global mean for
+    consistent logging on the main rank.
+
+    Args:
+        loss: Scalar tensor containing this rank's local weighted-mean loss.
+        loss_weight: Scalar tensor containing this rank's summed loss weights.
+        is_ddp: Whether the trainer module is wrapped in DDP.
+        process_group: Draft data-parallel process group, or ``None`` for world.
+
+    Returns:
+        Tuple of the scaled scalar loss for backward and the detached scalar
+        global weighted mean for logging.
+    """
+    if not is_ddp or not (dist.is_available() and dist.is_initialized()):
+        return loss, loss.detach()
+
+    weight = loss_weight.detach().to(device=loss.device, dtype=loss.dtype)
+    stats = torch.stack((loss.detach() * weight, weight))
+    dist.all_reduce(stats, op=dist.ReduceOp.SUM, group=process_group)
+    global_weight = stats[1].clamp_min(torch.finfo(stats.dtype).tiny)
+    dp_size = dist.get_world_size(group=process_group)
+    scaled_loss = loss * (weight * dp_size / global_weight)
+    return scaled_loss, stats[0] / global_weight
+
+
 def _submesh_or_none(device_mesh, name: str):
     """Return the named (flattened) submesh, or None if absent / no mesh.
 
@@ -616,6 +653,9 @@ def _build_trainer_module(self, attention_backend: str, recipe_cfg):
             block_size=self.block_size,
             attention_backend=attention_backend,
             num_anchors=int(recipe_cfg.get("num_anchors", 512)),
+            max_total_anchors=(
+                int(recipe_cfg["max_total_anchors"]) if recipe_cfg.get("max_total_anchors", None) is not None else None
+            ),
             # Paper default (Appendix A.1) for the shipped block_size=16 configs;
             # matches DFlashDecayLoss's own default. Set null explicitly in YAML
             # to disable the position decay (uniform weighting).
@@ -625,6 +665,8 @@ def _build_trainer_module(self, attention_backend: str, recipe_cfg):
             loss_type=str(recipe_cfg.get("loss_type", None) or "dflash"),
             prefix_weight_base=float(recipe_cfg.get("prefix_weight_base", 0.9)),
             sliding_window=self.draft_sliding_window,
+            use_fused_linear_ce=bool(recipe_cfg.get("use_fused_linear_ce", False)),
+            linear_ce_chunk_size=int(recipe_cfg.get("linear_ce_chunk_size", 1024)),
         )
 
     def _run_trainer_step(self, target_batch):
@@ -638,6 +680,14 @@ def _run_trainer_step(self, target_batch):
             doc_remaining=target_batch.doc_remaining,
         )
 
+    def _loss_terms(self, metrics) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
+        """Return differentiable local-mean loss terms and their denominators.
+
+        Composite trainers override this when their terms use different valid
+        sets and therefore need separate distributed normalization.
+        """
+        return ((metrics.loss, metrics.loss_weight),)
+
     def _log_extra_train_metrics(self, epoch_idx: int) -> None:
         """Hook for subclasses to log extra per-step metrics at a log point (no-op here)."""
 
@@ -1034,6 +1084,7 @@ def run_train_validation_loop(self):
                 last_batch_idx = -1
                 num_batches = len(self.train_dataloader)
                 is_ddp = isinstance(self.trainer_module, DistributedDataParallel)
+                draft_ddp_group = self._draft_ddp_process_group() if is_ddp else None
                 for batch_idx, batch in enumerate(self.train_dataloader):
                     last_batch_idx = batch_idx
                     batch = {k: v.to(self.device, non_blocking=True) for k, v in batch.items()}
@@ -1055,7 +1106,6 @@ def run_train_validation_loop(self):
                         local_has_valid = 1
                         try:
                             metrics = self._run_trainer_step(target_batch)
-                            loss = metrics.loss / self.grad_accumulation_steps
                         except NoValidAnchorsError:
                             # Every sample in this micro-batch is too short to form a
                             # block; nothing to learn from it on this rank.
@@ -1064,14 +1114,23 @@ def run_train_validation_loop(self):
                         # leaves one rank issuing its gradient all-reduce alone (DDP
                         # hang) or desyncs the accumulation windows: if ANY rank has no
                         # valid anchors, ALL ranks skip this micro-batch together.
-                        all_have_valid = _all_ranks_have_valid(local_has_valid, is_ddp, self.device)
+                        all_have_valid = _all_ranks_have_valid(
+                            local_has_valid, is_ddp, self.device, process_group=draft_ddp_group
+                        )
                         if all_have_valid:
+                            normalized_terms = [
+                                _normalize_ddp_loss(term, weight, is_ddp, process_group=draft_ddp_group)
+                                for term, weight in self._loss_terms(metrics)
+                            ]
+                            loss = sum(term for term, _ in normalized_terms)
+                            global_loss = sum(logged for _, logged in normalized_terms)
+                            loss = loss / self.grad_accumulation_steps
                             loss.backward()
                     if not all_have_valid:
                         self._skipped_micro_batches += 1
                         continue
 
-                    running_loss += metrics.loss.detach().item()
+                    running_loss += global_loss.item()
                     running_acc += metrics.accuracy.detach().item()
                     running_micro += 1
                     # Accumulated here, past the skip ``continue`` above, so a
@@ -1081,7 +1140,7 @@ def run_train_validation_loop(self):
                         totals = running_extra.setdefault(name, [0.0, 0.0])
                         totals[0] += numerator
                         totals[1] += denominator
-                    epoch_loss += metrics.loss.detach().item()
+                    epoch_loss += global_loss.item()
                     micro_step += 1
                     pending_micro_batches += 1
 
diff --git a/nemo_automodel/recipes/llm/train_dflash2.py b/nemo_automodel/recipes/llm/train_dflash2.py
index 280ba79f26..8be80a4869 100644
--- a/nemo_automodel/recipes/llm/train_dflash2.py
+++ b/nemo_automodel/recipes/llm/train_dflash2.py
@@ -73,6 +73,7 @@ def _build_trainer_module(self, attention_backend: str, recipe_cfg):
                 "loss_type is only supported by the DFlash recipe; the DFlash 2 trainer teacher-forces the "
                 "selector's predecessor from the fixed-anchor block layout and would silently ignore it."
             )
+        self.selector_loss_weight = float(recipe_cfg.get("selector_loss_weight", 1.0))
         return DFlash2TrainerModule(
             draft_model=self.draft_model,
             target_lm_head=self.target_model.get_output_embeddings(),
@@ -81,11 +82,14 @@ def _build_trainer_module(self, attention_backend: str, recipe_cfg):
             block_size=self.block_size,
             attention_backend=attention_backend,
             num_anchors=int(recipe_cfg.get("num_anchors", 512)),
+            max_total_anchors=(
+                int(recipe_cfg["max_total_anchors"]) if recipe_cfg.get("max_total_anchors", None) is not None else None
+            ),
             # Paper default (Appendix A.1) for the shipped block_size=16 configs;
             # matches DFlashDecayLoss's own default. Set null explicitly in YAML
             # to disable the position decay (uniform weighting).
             loss_decay_gamma=recipe_cfg.get("loss_decay_gamma", 7.0),
-            selector_loss_weight=float(recipe_cfg.get("selector_loss_weight", 1.0)),
+            selector_loss_weight=self.selector_loss_weight,
             sliding_window=self.draft_sliding_window,
         )
 
@@ -100,6 +104,21 @@ def _run_trainer_step(self, target_batch):
         self._last_dflash2_metrics = metrics
         return metrics
 
+    def _loss_terms(self, metrics) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
+        """Normalize backbone and selector losses over their own valid sets."""
+        weighted_selector_loss = self.selector_loss_weight * metrics.selector_loss
+        return (
+            (metrics.base_loss, metrics.loss_weight),
+            (weighted_selector_loss, metrics.selector_loss_denominator),
+        )
+
+    def _run_eval(self):
+        """Report the composite validation loss with per-term denominators."""
+        result = super()._run_eval()
+        if result is not None:
+            result["val_loss"] = result["val_base_loss"] + self.selector_loss_weight * result["val_selector_loss"]
+        return result
+
     def _log_extra_train_metrics(self, epoch_idx: int) -> None:
         """Log the DFlash 2 diagnostics for the most recent step (rank-0 local)."""
         m = getattr(self, "_last_dflash2_metrics", None)
@@ -143,11 +162,15 @@ def _extra_eval_metric_sums(self, metrics) -> dict[str, tuple[torch.Tensor, torc
         The shared validation loop SUM-reduces each pair before division.
         """
         loss_weight = metrics.loss_weight.detach()
+        selector_loss_denominator = metrics.selector_loss_denominator.detach()
         valid_tokens = metrics.valid_tokens.detach()
         valid_blocks = metrics.valid_blocks.detach()
         return {
             "val_base_loss": (metrics.base_loss.detach() * loss_weight, loss_weight),
-            "val_selector_loss": (metrics.selector_loss.detach() * loss_weight, loss_weight),
+            "val_selector_loss": (
+                metrics.selector_loss.detach() * selector_loss_denominator,
+                selector_loss_denominator,
+            ),
             "val_base_accuracy": (metrics.base_correct_tokens.detach(), valid_tokens),
             "val_base_accept_len": (metrics.base_accept_len_sum.detach(), valid_blocks),
             "val_candidate_recall": (metrics.candidate_recall.detach() * valid_tokens, valid_tokens),
diff --git a/nemo_automodel/recipes/llm/train_domino.py b/nemo_automodel/recipes/llm/train_domino.py
index bdb8e216dd..fd10cb7029 100644
--- a/nemo_automodel/recipes/llm/train_domino.py
+++ b/nemo_automodel/recipes/llm/train_domino.py
@@ -69,6 +69,9 @@ def _build_trainer_module(self, attention_backend: str, recipe_cfg):
             block_size=self.block_size,
             attention_backend=attention_backend,
             num_anchors=int(recipe_cfg.get("num_anchors", 512)),
+            max_total_anchors=(
+                int(recipe_cfg["max_total_anchors"]) if recipe_cfg.get("max_total_anchors", None) is not None else None
+            ),
             # Paper default (Appendix A.1) for the shipped block_size=16 configs;
             # matches DFlashDecayLoss's own default. Set null explicitly in YAML
             # to disable the position decay (uniform weighting).
diff --git a/nemo_automodel/recipes/llm/train_jetspec.py b/nemo_automodel/recipes/llm/train_jetspec.py
index f1b8e67c10..dd2803d62e 100644
--- a/nemo_automodel/recipes/llm/train_jetspec.py
+++ b/nemo_automodel/recipes/llm/train_jetspec.py
@@ -77,6 +77,9 @@ def _build_trainer_module(self, attention_backend: str, recipe_cfg):
             block_size=self.block_size,
             attention_backend=attention_backend,
             num_anchors=int(recipe_cfg.get("num_anchors", 512)),
+            max_total_anchors=(
+                int(recipe_cfg["max_total_anchors"]) if recipe_cfg.get("max_total_anchors", None) is not None else None
+            ),
             kd_temperature=float(recipe_cfg.get("kd_temperature", 1.0)),
             kd_chunk_size=int(recipe_cfg.get("kd_chunk_size", 0)),
         )
diff --git a/tests/unit_tests/loss/test_dllm_loss.py b/tests/unit_tests/loss/test_dllm_loss.py
index c596edcfd3..05fec41532 100644
--- a/tests/unit_tests/loss/test_dllm_loss.py
+++ b/tests/unit_tests/loss/test_dllm_loss.py
@@ -397,7 +397,7 @@ def test_fused_matches_nonfused(self):
 
         logits = torch.nn.functional.linear(hidden, weight, bias)
         ref = loss_fn(logits, target_ids, block_mask, num_tokens=B * T, block_size=bs)
-        fused = loss_fn.forward_fused(
+        fused, fused_correct = loss_fn.forward_fused_with_correct(
             hidden,
             weight,
             target_ids,
@@ -410,6 +410,33 @@ def test_fused_matches_nonfused(self):
         assert torch.allclose(ref.total_loss, fused.total_loss, atol=1e-4)
         assert torch.equal(ref.draft_correct_per_pos, fused.draft_correct_per_pos)
         assert torch.equal(ref.draft_count_per_pos, fused.draft_count_per_pos)
+        assert torch.equal(fused_correct, logits.argmax(dim=-1) == target_ids)
+        assert len(loss_fn.forward_fused(hidden, weight, target_ids, block_mask, block_size=bs)) == 4
+
+    def test_fused_detailed_projection_matches_nonfused(self):
+        """Detailed fused projection returns correctness without expanding the shared output."""
+        torch.manual_seed(4)
+        B, N, bs, D, V = 2, 2, 4, 16, 32
+        T = N * (bs - 1)
+        hidden = torch.randn(B, T, D, requires_grad=True)
+        lm_head = torch.nn.Linear(D, V)
+        target_ids = torch.randint(0, V, (B, T))
+        block_mask = torch.ones(B, T)
+        loss_fn = DFlashDecayLoss(loss_gamma=4.0, use_fused_linear_ce=True, chunk_size=3)
+
+        ref = loss_fn(lm_head(hidden), target_ids, block_mask, num_tokens=B * T, block_size=bs)
+        fused, fused_correct = loss_fn.forward_fused_with_correct(
+            hidden,
+            lm_head.weight,
+            target_ids,
+            block_mask,
+            num_tokens=B * T,
+            block_size=bs,
+            lm_head_bias=lm_head.bias,
+        )
+
+        torch.testing.assert_close(fused.total_loss, ref.total_loss, atol=1e-5, rtol=1e-5)
+        assert torch.equal(fused_correct, lm_head(hidden).argmax(dim=-1) == target_ids)
 
     def test_paper_default_first_offset_weight_is_one(self):
         """The first predicted position of every block must have decay weight 1.0
@@ -426,12 +453,12 @@ def test_paper_default_first_offset_weight_is_one(self):
             # First weight of every block must be 1.0; weights must decay within a block.
             for b in range(n_blocks):
                 start = b * T_per
-                assert torch.isclose(w[start], torch.tensor(1.0)), (
-                    f"block_size={block_size}, block {b}: first weight {w[start].item()} != 1.0"
-                )
-                assert w[start] > w[start + T_per - 1], (
-                    f"block_size={block_size}, block {b}: weights do not decay within block"
-                )
+                assert torch.isclose(
+                    w[start], torch.tensor(1.0)
+                ), f"block_size={block_size}, block {b}: first weight {w[start].item()} != 1.0"
+                assert (
+                    w[start] > w[start + T_per - 1]
+                ), f"block_size={block_size}, block {b}: weights do not decay within block"
 
     def test_recipe_per_pos_metrics_dict_construction(self):
         """Lock the recipe-side contract: given the loss's per-rank
diff --git a/tests/unit_tests/recipes/llm/test_train_dflash.py b/tests/unit_tests/recipes/llm/test_train_dflash.py
index 6a4fba0ed9..6a2ce37f7c 100644
--- a/tests/unit_tests/recipes/llm/test_train_dflash.py
+++ b/tests/unit_tests/recipes/llm/test_train_dflash.py
@@ -21,10 +21,13 @@
 from __future__ import annotations
 
 import logging
+from contextlib import nullcontext
+from datetime import timedelta
 from types import SimpleNamespace
 
 import pytest
 import torch
+import torch.multiprocessing as mp
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
 
 from nemo_automodel.components.checkpoint.config import CheckpointingConfig
@@ -40,6 +43,109 @@
 _TARGET_LAYER_IDS = [1, 3, 5]
 
 
+_DDP_BASE_BATCHES = (
+    (
+        ([[1.0, 0.0], [0.0, 1.0]], [1.0, -1.0]),
+        ([[2.0, 1.0]], [0.5]),
+    ),
+    (
+        ([[1.0, 1.0]], [0.0]),
+        ([[0.0, 2.0], [2.0, 0.0], [1.0, -1.0]], [1.5, -0.5, 0.25]),
+    ),
+)
+_DDP_SELECTOR_BATCHES = (
+    (
+        ([[0.5, 1.0]], [0.25]),
+        ([[1.5, -0.5], [0.25, 2.0]], [0.75, -1.0]),
+    ),
+    (
+        ([[1.0, 0.5], [-1.0, 1.0]], [0.0, 1.25]),
+        ([], []),
+    ),
+)
+
+
+def _weighted_mse_term(prediction: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return a local mean and count, preserving a zero gradient for an empty set.
+
+    Args:
+        prediction: Tensor of shape ``[rows, 1]`` containing model outputs.
+        target: Tensor of shape ``[rows]`` containing regression targets.
+
+    Returns:
+        Scalar local-mean loss and scalar row-count tensors.
+    """
+    count = torch.tensor(float(target.numel()), dtype=prediction.dtype)
+    if target.numel() == 0:
+        return prediction.sum() * 0.0, count
+    return (prediction.squeeze(-1) - target).square().mean(), count
+
+
+def _ddp_batch(spec, rank: int, micro_step: int) -> tuple[torch.Tensor, torch.Tensor]:
+    inputs, targets = spec[rank][micro_step]
+    return torch.tensor(inputs, dtype=torch.float32).reshape(-1, 2), torch.tensor(targets, dtype=torch.float32)
+
+
+def _run_two_rank_weighted_loss_gradient_parity(rank: int, init_file: str) -> None:
+    """Exercise real DDP/no_sync with uneven base and selector denominators."""
+    import torch.distributed as dist
+
+    dist.init_process_group(
+        "gloo",
+        init_method=f"file://{init_file}",
+        rank=rank,
+        world_size=2,
+        timeout=timedelta(seconds=60),
+    )
+    try:
+        torch.manual_seed(17)
+        module = torch.nn.Linear(2, 1)
+        ddp = torch.nn.parallel.DistributedDataParallel(module)
+        initial_state = {name: value.detach().clone() for name, value in ddp.module.state_dict().items()}
+        selector_coefficient = 0.3
+
+        for micro_step in range(2):
+            base_x, base_y = _ddp_batch(_DDP_BASE_BATCHES, rank, micro_step)
+            selector_x, selector_y = _ddp_batch(_DDP_SELECTOR_BATCHES, rank, micro_step)
+            all_x = torch.cat((base_x, selector_x))
+            sync_context = ddp.no_sync() if micro_step == 0 else nullcontext()
+            with sync_context:
+                prediction = ddp(all_x)
+                base_prediction = prediction[: base_x.shape[0]]
+                selector_prediction = prediction[base_x.shape[0] :]
+                base_loss, base_weight = _weighted_mse_term(base_prediction, base_y)
+                selector_loss, selector_weight = _weighted_mse_term(selector_prediction, selector_y)
+                normalized_base, _ = train_dflash._normalize_ddp_loss(base_loss, base_weight, True)
+                normalized_selector, _ = train_dflash._normalize_ddp_loss(
+                    selector_coefficient * selector_loss, selector_weight, True
+                )
+                ((normalized_base + normalized_selector) / 2).backward()
+
+        reference = torch.nn.Linear(2, 1)
+        reference.load_state_dict(initial_state)
+        reference_loss = torch.zeros(())
+        for micro_step in range(2):
+            for spec, coefficient in (
+                (_DDP_BASE_BATCHES, 1.0),
+                (_DDP_SELECTOR_BATCHES, selector_coefficient),
+            ):
+                xs, ys = zip(*(_ddp_batch(spec, other_rank, micro_step) for other_rank in range(2)))
+                all_x = torch.cat(xs)
+                all_y = torch.cat(ys)
+                if all_y.numel() > 0:
+                    reference_loss = (
+                        reference_loss + coefficient * (reference(all_x).squeeze(-1) - all_y).square().mean() / 2
+                    )
+        reference_loss.backward()
+
+        for (name, parameter), reference_parameter in zip(ddp.module.named_parameters(), reference.parameters()):
+            torch.testing.assert_close(
+                parameter.grad, reference_parameter.grad, msg=lambda message: f"{name}: {message}"
+            )
+    finally:
+        dist.destroy_process_group()
+
+
 def _dflash_draft():
     cfg = Qwen3Config(
         vocab_size=_VOCAB,
@@ -75,6 +181,27 @@ def _bare_dflash_recipe():
     return recipe
 
 
+def test_build_trainer_module_wires_fused_ce_and_anchor_budget():
+    recipe = _bare_dflash_recipe()
+
+    module = recipe._build_trainer_module(
+        "sdpa",
+        {
+            "num_anchors": 512,
+            "max_total_anchors": 512,
+            "use_fused_linear_ce": True,
+            "linear_ce_chunk_size": 256,
+            "loss_decay_gamma": 4.0,
+        },
+    )
+
+    assert module.num_anchors == 512
+    assert module.max_total_anchors == 512
+    assert module.use_fused_linear_ce is True
+    assert module.loss_fn.chunk_size == 256
+    assert module.loss_decay_gamma == 4.0
+
+
 def _ckpt_self(ckpt_every_steps, save_every_epoch, global_step, total_optim_steps=None):
     calls = []
     return (
@@ -392,14 +519,54 @@ def test_all_ranks_have_valid_ddp_min_reduces_across_ranks(monkeypatch):
     monkeypatch.setattr(train_dflash.dist, "is_initialized", lambda: True)
 
     # Another rank skipped -> the collective drives the flag to 0 -> all skip.
-    monkeypatch.setattr(train_dflash.dist, "all_reduce", lambda t, op=None: t.fill_(0))
+    monkeypatch.setattr(train_dflash.dist, "all_reduce", lambda t, op=None, group=None: t.fill_(0))
     assert train_dflash._all_ranks_have_valid(1, is_ddp=True, device="cpu") is False
 
     # Every rank valid -> MIN leaves the 1 in place -> all run the backward.
-    monkeypatch.setattr(train_dflash.dist, "all_reduce", lambda t, op=None: None)
+    monkeypatch.setattr(train_dflash.dist, "all_reduce", lambda t, op=None, group=None: None)
     assert train_dflash._all_ranks_have_valid(1, is_ddp=True, device="cpu") is True
 
 
+def test_normalize_ddp_loss_scales_gradient_by_global_weight(monkeypatch):
+    monkeypatch.setattr(train_dflash.dist, "is_available", lambda: True)
+    monkeypatch.setattr(train_dflash.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(train_dflash.dist, "get_world_size", lambda group=None: 2)
+
+    def fake_all_reduce(stats, op=None, group=None):
+        # Local numerator/weight are 6/2; the other rank contributes 20/8.
+        stats.add_(torch.tensor([20.0, 8.0]))
+
+    monkeypatch.setattr(train_dflash.dist, "all_reduce", fake_all_reduce)
+    loss = torch.tensor(3.0, requires_grad=True)
+
+    scaled_loss, global_loss = train_dflash._normalize_ddp_loss(
+        loss, torch.tensor(2.0), is_ddp=True, process_group="draft-dp"
+    )
+    scaled_loss.backward()
+
+    assert loss.grad.item() == pytest.approx(0.4)
+    assert global_loss.item() == pytest.approx(2.6)
+
+
+def test_normalize_ddp_loss_is_noop_without_ddp():
+    loss = torch.tensor(3.0, requires_grad=True)
+
+    scaled_loss, logged_loss = train_dflash._normalize_ddp_loss(loss, torch.tensor(2.0), is_ddp=False)
+
+    assert scaled_loss is loss
+    assert logged_loss.item() == 3.0
+
+
+def test_two_rank_ddp_weighted_loss_matches_global_reference(tmp_path):
+    """Real DDP accumulation matches global per-term means with uneven rank work."""
+    mp.spawn(
+        _run_two_rank_weighted_loss_gradient_parity,
+        args=(str(tmp_path / "dflash_weighted_loss"),),
+        nprocs=2,
+        join=True,
+    )
+
+
 def test_all_reduce_sum_uses_additive_distributed_statistics(monkeypatch):
     monkeypatch.setattr(train_dflash.dist, "is_available", lambda: True)
     monkeypatch.setattr(train_dflash.dist, "is_initialized", lambda: True)
diff --git a/tests/unit_tests/recipes/llm/test_train_dflash2.py b/tests/unit_tests/recipes/llm/test_train_dflash2.py
index 3bf036bbfc..d292699fba 100644
--- a/tests/unit_tests/recipes/llm/test_train_dflash2.py
+++ b/tests/unit_tests/recipes/llm/test_train_dflash2.py
@@ -33,6 +33,7 @@
 )
 from nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2 import Qwen3DFlash2DraftModel
 from nemo_automodel.components.speculative.dflash.registry import resolve_dflash_draft_spec
+from nemo_automodel.recipes.llm.train_dflash import TrainDFlashRecipe
 from nemo_automodel.recipes.llm.train_dflash2 import TrainDFlash2Recipe
 
 VOCAB = 64
@@ -128,10 +129,17 @@ def test_build_trainer_module_is_dflash2():
     recipe.target_model = _target_model()
     recipe.draft_sliding_window = 2048
     module = recipe._build_trainer_module(
-        "sdpa", {"num_anchors": 7, "loss_decay_gamma": 5.0, "selector_loss_weight": 0.25}
+        "sdpa",
+        {
+            "num_anchors": 7,
+            "max_total_anchors": 11,
+            "loss_decay_gamma": 5.0,
+            "selector_loss_weight": 0.25,
+        },
     )
     assert isinstance(module, DFlash2TrainerModule)
     assert module.num_anchors == 7
+    assert module.max_total_anchors == 11
     assert module.loss_decay_gamma == 5.0
     assert module.selector_loss_weight == 0.25
     # The window the recipe resolved must reach the trainer, or training and the
@@ -176,6 +184,7 @@ def _metrics(**overrides):
         valid_blocks=torch.tensor(2.0),
         base_loss=torch.tensor(2.3),
         selector_loss=torch.tensor(0.7),
+        selector_loss_denominator=torch.tensor(6.0),
         base_accuracy=torch.tensor(0.4),
         base_correct_tokens=torch.tensor(4.0),
         base_accept_len=torch.tensor(4.0),
@@ -227,7 +236,8 @@ def test_eval_sums_keep_additive_numerators():
 
     assert sums["val_base_loss"][0].item() == pytest.approx(18.4)
     assert sums["val_base_loss"][1].item() == 8.0
-    assert sums["val_selector_loss"][0].item() == pytest.approx(5.6)
+    assert sums["val_selector_loss"][0].item() == pytest.approx(4.2)
+    assert sums["val_selector_loss"][1].item() == 6.0
     assert sums["val_base_accuracy"][0].item() == 4.0
     assert sums["val_base_accept_len"][0].item() == 8.0
     assert sums["val_base_accept_len"][1].item() == 2.0
@@ -266,6 +276,35 @@ def _fake_module(**kwargs):
     assert recipe._last_dflash2_metrics is out
 
 
+def test_loss_terms_keep_base_and_selector_denominators_separate():
+    recipe = _recipe()
+    recipe.selector_loss_weight = 0.5
+    metrics = _metrics(
+        loss=torch.tensor(3.0, requires_grad=True),
+        base_loss=torch.tensor(2.0, requires_grad=True),
+        loss_weight=torch.tensor(8.0),
+        selector_loss_denominator=torch.tensor(3.0),
+    )
+
+    terms = recipe._loss_terms(metrics)
+
+    assert terms[0][0] is metrics.base_loss
+    assert terms[0][1].item() == 8.0
+    assert terms[1][0].item() == pytest.approx(0.35)
+    assert terms[1][1].item() == 3.0
+
+
+def test_run_eval_rebuilds_composite_loss_from_separately_weighted_terms(monkeypatch):
+    recipe = _recipe()
+    recipe.selector_loss_weight = 0.25
+    base_result = {"val_loss": 99.0, "val_base_loss": 2.0, "val_selector_loss": 4.0}
+    monkeypatch.setattr(TrainDFlashRecipe, "_run_eval", lambda self: dict(base_result))
+
+    result = recipe._run_eval()
+
+    assert result["val_loss"] == 3.0
+
+
 def test_setup_resets_the_metrics_cache(monkeypatch):
     recipe = _recipe()
     # Bypass the heavy DFlash setup (super().setup()); only the reset is under test.
diff --git a/tests/unit_tests/recipes/llm/test_train_dflash_noanchor_skip.py b/tests/unit_tests/recipes/llm/test_train_dflash_noanchor_skip.py
index 609c59e380..9623ce6898 100644
--- a/tests/unit_tests/recipes/llm/test_train_dflash_noanchor_skip.py
+++ b/tests/unit_tests/recipes/llm/test_train_dflash_noanchor_skip.py
@@ -53,6 +53,7 @@ def forward(self, **kwargs):
         # accumulates them per micro-batch to average over the log window.
         return SimpleNamespace(
             loss=out.abs() + 1.0,
+            loss_weight=torch.tensor(4.0),
             accuracy=torch.tensor(0.5),
             accept_len_sum=torch.tensor(2.0),
             valid_blocks=torch.tensor(1.0),
diff --git a/tests/unit_tests/recipes/llm/test_train_domino.py b/tests/unit_tests/recipes/llm/test_train_domino.py
index e6a904f1c3..804c904871 100644
--- a/tests/unit_tests/recipes/llm/test_train_domino.py
+++ b/tests/unit_tests/recipes/llm/test_train_domino.py
@@ -108,9 +108,10 @@ def test_build_trainer_module_is_domino():
         get_output_embeddings=lambda: torch.nn.Linear(HIDDEN, VOCAB, bias=False),
         get_input_embeddings=lambda: torch.nn.Embedding(VOCAB, HIDDEN),
     )
-    module = recipe._build_trainer_module("sdpa", {"num_anchors": 7, "loss_decay_gamma": 5.0})
+    module = recipe._build_trainer_module("sdpa", {"num_anchors": 7, "max_total_anchors": 11, "loss_decay_gamma": 5.0})
     assert isinstance(module, DominoTrainerModule)
     assert module.num_anchors == 7
+    assert module.max_total_anchors == 11
     assert module.loss_decay_gamma == 5.0
     assert module.shift_label is True
 
diff --git a/tests/unit_tests/recipes/llm/test_train_jetspec.py b/tests/unit_tests/recipes/llm/test_train_jetspec.py
index ed31b5f070..460747754b 100644
--- a/tests/unit_tests/recipes/llm/test_train_jetspec.py
+++ b/tests/unit_tests/recipes/llm/test_train_jetspec.py
@@ -113,9 +113,13 @@ def test_build_dflash_config_stamps_causal():
 
 def test_build_trainer_module_is_jetspec():
     recipe = _jetspec_recipe()
-    module = recipe._build_trainer_module("sdpa", {"num_anchors": 7, "kd_temperature": 2.0, "kd_chunk_size": 64})
+    module = recipe._build_trainer_module(
+        "sdpa",
+        {"num_anchors": 7, "max_total_anchors": 11, "kd_temperature": 2.0, "kd_chunk_size": 64},
+    )
     assert isinstance(module, JetSpecTrainerModule)
     assert module.num_anchors == 7
+    assert module.max_total_anchors == 11
     assert module.kd_temperature == 2.0
     assert module.kd_chunk_size == 64
 
diff --git a/tests/unit_tests/speculative/test_dflash2_draft.py b/tests/unit_tests/speculative/test_dflash2_draft.py
index 8b5d268f4c..a0aacbd5fa 100644
--- a/tests/unit_tests/speculative/test_dflash2_draft.py
+++ b/tests/unit_tests/speculative/test_dflash2_draft.py
@@ -406,6 +406,7 @@ def __init__(self, cfg: Qwen3Config, forced_token_id: int):
         super().__init__()
         self.model = torch.nn.Module()
         self.model.embed_tokens = torch.nn.Embedding(cfg.vocab_size, cfg.hidden_size)
+        self.model.layers = torch.nn.ModuleList([torch.nn.Identity() for _ in range(cfg.num_target_layers)])
         self.lm_head = torch.nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
         self.num_layers = cfg.num_target_layers
         self.vocab_size = cfg.vocab_size
@@ -422,6 +423,8 @@ def forward(
         output_hidden_states=False,
     ):
         hidden = self.model.embed_tokens(input_ids)
+        for layer in self.model.layers:
+            hidden = layer(hidden)
         keep = input_ids.shape[1] if logits_to_keep is None else logits_to_keep
         # A very peaked distribution so the sampled path is deterministic too.
         logits = torch.zeros(input_ids.shape[0], keep, self.vocab_size)
diff --git a/tests/unit_tests/speculative/test_dflash_core.py b/tests/unit_tests/speculative/test_dflash_core.py
index 7debbd89c0..3c09bb3dd0 100644
--- a/tests/unit_tests/speculative/test_dflash_core.py
+++ b/tests/unit_tests/speculative/test_dflash_core.py
@@ -35,7 +35,15 @@
 MASK_ID = VOCAB - 1
 
 
-def _build_trainer(num_anchors=8, loss_decay_gamma=None, attention_backend="sdpa", sliding_window=None):
+def _build_trainer(
+    num_anchors=8,
+    loss_decay_gamma=None,
+    attention_backend="sdpa",
+    sliding_window=None,
+    max_total_anchors=None,
+    use_fused_linear_ce=False,
+    linear_ce_chunk_size=1024,
+):
     cfg = Qwen3Config(
         vocab_size=VOCAB,
         hidden_size=HIDDEN,
@@ -56,6 +64,8 @@ def _build_trainer(num_anchors=8, loss_decay_gamma=None, attention_backend="sdpa
     draft = Qwen3DFlashDraftModel(cfg)
     lm_head = torch.nn.Linear(HIDDEN, VOCAB, bias=False)
     embed = torch.nn.Embedding(VOCAB, HIDDEN)
+    lm_head.requires_grad_(False)
+    embed.requires_grad_(False)
     return DFlashTrainerModule(
         draft_model=draft,
         target_lm_head=lm_head,
@@ -64,8 +74,11 @@ def _build_trainer(num_anchors=8, loss_decay_gamma=None, attention_backend="sdpa
         block_size=BLOCK_SIZE,
         attention_backend=attention_backend,
         num_anchors=num_anchors,
+        max_total_anchors=max_total_anchors,
         loss_decay_gamma=loss_decay_gamma,
         sliding_window=sliding_window,
+        use_fused_linear_ce=use_fused_linear_ce,
+        linear_ce_chunk_size=linear_ce_chunk_size,
     )
 
 
@@ -151,6 +164,28 @@ def test_sample_anchor_positions_respect_loss_mask():
     assert (valid_anchors < 10).all()
 
 
+def test_max_total_anchors_caps_constructed_batch_slots():
+    trainer = _build_trainer(num_anchors=8, max_total_anchors=5)
+    loss_mask = torch.ones(3, 20)
+
+    anchors, keep = trainer._sample_anchor_positions(20, loss_mask, torch.device("cpu"))
+
+    # Anchor tensors are rectangular, so the cap is divided uniformly across
+    # the local micro-batch. The unused remainder is intentional: it keeps the
+    # actual draft-block allocation, not only the valid-count sum, under budget.
+    assert anchors.shape == (3, 1)
+    assert anchors.numel() <= 5
+    assert keep.sum().item() == 3
+
+
+def test_max_total_anchors_must_cover_local_batch():
+    trainer = _build_trainer(num_anchors=8, max_total_anchors=1)
+    loss_mask = torch.ones(2, 20)
+
+    with pytest.raises(ValueError, match="max_total_anchors"):
+        trainer._sample_anchor_positions(20, loss_mask, torch.device("cpu"))
+
+
 def test_no_valid_anchors_raises():
     trainer = _build_trainer()
     seq_len = 20
@@ -202,6 +237,76 @@ def test_invalid_prefix_weight_base_raises():
         )
 
 
+def test_fused_linear_ce_rejects_variable_prefix_objective():
+    trainer = _build_trainer()
+    with pytest.raises(ValueError, match="use_fused_linear_ce"):
+        DFlashTrainerModule(
+            draft_model=trainer.draft_model,
+            target_lm_head=trainer.lm_head,
+            target_embed_tokens=trainer.embed_tokens,
+            mask_token_id=MASK_ID,
+            block_size=BLOCK_SIZE,
+            loss_type="variable_prefix",
+            use_fused_linear_ce=True,
+        )
+
+
+def test_constructor_preserves_historical_positional_prefix():
+    trainer = _build_trainer()
+
+    positional = DFlashTrainerModule(
+        trainer.draft_model,
+        trainer.lm_head,
+        trainer.embed_tokens,
+        MASK_ID,
+        BLOCK_SIZE,
+        "sdpa",
+        8,
+        4.0,
+        "dflash",
+        0.9,
+    )
+
+    assert positional.loss_decay_gamma == 4.0
+    assert positional.loss_type == "dflash"
+    assert positional.prefix_weight_base == 0.9
+    assert positional.max_total_anchors is None
+
+
+def test_fused_linear_ce_requires_frozen_target_head():
+    trainer = _build_trainer(use_fused_linear_ce=True)
+    trainer.lm_head.requires_grad_(True)
+    input_ids, hidden, loss_mask = _inputs()
+
+    with pytest.raises(ValueError, match="frozen target LM head"):
+        trainer(input_ids=input_ids, hidden_states=hidden, loss_mask=loss_mask)
+
+
+def test_fused_linear_ce_matches_dense_metrics_and_draft_gradients():
+    trainer = _build_trainer(loss_decay_gamma=4.0, use_fused_linear_ce=False, linear_ce_chunk_size=3)
+    input_ids, hidden, loss_mask = _inputs(bsz=2, seq_len=16)
+
+    torch.manual_seed(11)
+    dense = trainer(input_ids=input_ids, hidden_states=hidden, loss_mask=loss_mask)
+    dense.loss.backward()
+    dense_grads = [p.grad.detach().clone() for p in trainer.draft_model.parameters() if p.grad is not None]
+
+    trainer.zero_grad(set_to_none=True)
+    trainer.use_fused_linear_ce = True
+    torch.manual_seed(11)
+    fused = trainer(input_ids=input_ids, hidden_states=hidden, loss_mask=loss_mask)
+    fused.loss.backward()
+    fused_grads = [p.grad.detach().clone() for p in trainer.draft_model.parameters() if p.grad is not None]
+
+    torch.testing.assert_close(fused.loss, dense.loss, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(fused.accuracy, dense.accuracy)
+    torch.testing.assert_close(fused.accept_len, dense.accept_len)
+    torch.testing.assert_close(fused.correct_tokens, dense.correct_tokens)
+    assert len(fused_grads) == len(dense_grads)
+    for fused_grad, dense_grad in zip(fused_grads, dense_grads):
+        torch.testing.assert_close(fused_grad, dense_grad, atol=1e-5, rtol=1e-4)
+
+
 def test_variable_prefix_forward_finite_loss_and_grads():
     trainer = _build_vp_trainer(loss_decay_gamma=7.0)
     input_ids, hidden, loss_mask = _inputs()
diff --git a/tests/unit_tests/speculative/test_dflash_draft.py b/tests/unit_tests/speculative/test_dflash_draft.py
index cd95b3403f..aa52a8cb70 100644
--- a/tests/unit_tests/speculative/test_dflash_draft.py
+++ b/tests/unit_tests/speculative/test_dflash_draft.py
@@ -18,14 +18,17 @@
 
 from types import SimpleNamespace
 
+import pytest
 import torch
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
 
+from nemo_automodel.components.speculative.dflash import draft_qwen3
 from nemo_automodel.components.speculative.dflash.draft_qwen3 import (
     Qwen3DFlashDraftModel,
     _sliding_window_mask,
     build_target_layer_ids,
     extract_context_feature,
+    forward_target_with_context_feature,
 )
 
 
@@ -49,6 +52,59 @@ def test_extract_context_feature_uses_offset_one():
     assert torch.allclose(out[..., 3:], torch.full((1, 2, 3), 4.0))
 
 
+def test_extract_context_feature_rejects_post_norm_final_entry():
+    hs = [torch.full((1, 2, 3), float(i)) for i in range(6)]
+
+    with pytest.raises(ValueError, match="cannot be recovered pre-norm"):
+        extract_context_feature(hs, [4])
+
+
+class _AddConstant(torch.nn.Module):
+    """Test layer that makes block and final-norm outputs distinguishable."""
+
+    def __init__(self, value: float) -> None:
+        super().__init__()
+        self.value = value
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """Add the configured scalar without changing tensor layout.
+
+        Args:
+            inputs: Tensor with arbitrary shape.
+
+        Returns:
+            Tensor with the same shape as ``inputs``.
+        """
+        return inputs + self.value
+
+
+def test_forward_target_with_context_feature_captures_final_block_before_model_norm():
+    cfg = _draft_cfg()
+    target = _ConstantTarget(cfg, forced_token_id=1)
+    target.model.layers = torch.nn.ModuleList([_AddConstant(float(i + 1)) for i in range(target.num_layers)])
+    target.model.norm = _AddConstant(10.0)
+    input_ids = torch.tensor([[1, 2, 3]])
+
+    output, feature = forward_target_with_context_feature(target, input_ids, [target.num_layers - 1])
+
+    expected = target.model.embed_tokens(input_ids)
+    for i in range(target.num_layers):
+        expected = expected + float(i + 1)
+    assert torch.equal(feature, expected)
+    assert torch.equal(output.final_hidden_state, target.model.norm(expected))
+
+
+def test_forward_target_with_context_feature_rejects_invalid_id_without_leaking_hook():
+    cfg = _draft_cfg()
+    target = _ConstantTarget(cfg, forced_token_id=1)
+    first_layer = target.model.layers[0]
+
+    with pytest.raises(ValueError, match="out of bounds"):
+        forward_target_with_context_feature(target, torch.tensor([[1]]), [0, target.num_layers])
+
+    assert not first_layer._forward_hooks
+
+
 def _draft_cfg(bs=4):
     cfg = Qwen3Config(
         vocab_size=64,
@@ -130,6 +186,52 @@ def test_draft_forward_output_shape():
     assert torch.isfinite(out).all()
 
 
+def test_flex_attention_forces_regular_kernel_for_short_dynamic_queries(monkeypatch):
+    cfg = _draft_cfg()
+    cfg._attn_implementation = "flex_attention"
+    captured_options = []
+
+    def fake_flex_attention(_module, query, _key, _value, _mask, **kwargs):
+        """Record kernel options while preserving the attention output contract.
+
+        Args:
+            _module: Attention module (unused by this test double).
+            query: Tensor of shape ``[batch, heads, query, head_dim]``.
+            _key: Tensor of shape ``[batch, heads, key_value, head_dim]``.
+            _value: Tensor of shape ``[batch, heads, key_value, head_dim]``.
+            _mask: Flex-attention block mask, or ``None``.
+            **kwargs: Attention options, including ``kernel_options``.
+
+        Returns:
+            Attention output tensor of shape ``[batch, query, heads, head_dim]``
+            and ``None`` for attention weights.
+        """
+        captured_options.append(kwargs["kernel_options"])
+        return query.transpose(1, 2), None
+
+    monkeypatch.setattr(
+        draft_qwen3,
+        "ALL_ATTENTION_FUNCTIONS",
+        {"flex_attention": fake_flex_attention},
+    )
+    draft = Qwen3DFlashDraftModel(cfg)
+    batch, context, query = 2, 10, 12
+    noise = torch.randn(batch, query, cfg.hidden_size)
+    target_hidden = torch.randn(batch, context, len(cfg.dflash_config["target_layer_ids"]) * cfg.hidden_size)
+    position_ids = torch.arange(context + query).unsqueeze(0).expand(batch, -1)
+
+    output = draft(
+        position_ids=position_ids,
+        attention_mask=None,
+        noise_embedding=noise,
+        target_hidden=target_hidden,
+    )
+
+    assert output.shape == noise.shape
+    assert len(captured_options) == cfg.num_hidden_layers
+    assert all(options["FORCE_USE_FLEX_ATTENTION"] for options in captured_options)
+
+
 class _ConstantTarget(torch.nn.Module):
     """Minimal stand-in for the verifier: always samples ``forced_token_id``.
 
@@ -141,6 +243,8 @@ def __init__(self, cfg: Qwen3Config, forced_token_id: int):
         super().__init__()
         self.model = torch.nn.Module()
         self.model.embed_tokens = torch.nn.Embedding(cfg.vocab_size, cfg.hidden_size)
+        self.model.layers = torch.nn.ModuleList([torch.nn.Identity() for _ in range(cfg.num_target_layers)])
+        self.model.norm = torch.nn.Identity()
         self.lm_head = torch.nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
         self.num_layers = cfg.num_target_layers
         self.vocab_size = cfg.vocab_size
@@ -157,10 +261,13 @@ def forward(
         output_hidden_states=False,
     ):
         hidden = self.model.embed_tokens(input_ids)
+        for layer in self.model.layers:
+            hidden = layer(hidden)
+        final_hidden = self.model.norm(hidden)
         keep = input_ids.shape[1] if logits_to_keep is None else logits_to_keep
         logits = torch.zeros(input_ids.shape[0], keep, self.vocab_size)
         logits[..., self.forced_token_id] = 1.0
-        return SimpleNamespace(logits=logits, hidden_states=[hidden] * (self.num_layers + 1))
+        return SimpleNamespace(logits=logits, final_hidden_state=final_hidden)
 
 
 def test_spec_generate_keeps_generated_tokens_equal_to_the_mask_id():
diff --git a/tests/unit_tests/speculative/test_dflash_target.py b/tests/unit_tests/speculative/test_dflash_target.py
index 3e23fa585a..306299022b 100644
--- a/tests/unit_tests/speculative/test_dflash_target.py
+++ b/tests/unit_tests/speculative/test_dflash_target.py
@@ -36,12 +36,30 @@
 _LAYERS = 4
 
 
+class _AddConstant(nn.Module):
+    def __init__(self, value: float) -> None:
+        super().__init__()
+        self.value = value
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """Add a scalar without changing tensor layout.
+
+        Args:
+            inputs: Tensor with arbitrary shape.
+
+        Returns:
+            Tensor with the same shape as ``inputs``.
+        """
+        return inputs + self.value
+
+
 class _FakeHFBackbone(nn.Module):
     """HuggingFace-style backbone: ``ModuleList`` layers, explicit HF flags."""
 
     def __init__(self, embed: nn.Embedding) -> None:
         super().__init__()
         self.layers = nn.ModuleList([nn.Linear(_HIDDEN, _HIDDEN) for _ in range(_LAYERS)])
+        self.norm = _AddConstant(10.0)
         self._embed = embed
 
     def forward(
@@ -50,7 +68,7 @@ def forward(
         h = self._embed(input_ids)
         for layer in self.layers:
             h = layer(h)
-        return (h,)
+        return (self.norm(h),)
 
 
 class _FakeHFCausalLM(nn.Module):
@@ -162,6 +180,21 @@ def test_generate_batch_concatenates_selected_layers():
     assert torch.equal(out.loss_mask, loss)
 
 
+def test_generate_batch_captures_final_block_before_model_norm():
+    model = _FakeHFCausalLM()
+    model.model.layers = nn.ModuleList([_AddConstant(float(index + 1)) for index in range(_LAYERS)])
+    target = HFDFlashTargetModel(model, target_layer_ids=[_LAYERS - 1])
+    input_ids, attention_mask, loss_mask = _batch(batch=2, seq=8)
+
+    out = target.generate_batch(input_ids, attention_mask, loss_mask)
+
+    expected_pre_norm = model.get_input_embeddings()(input_ids)
+    for index in range(_LAYERS):
+        expected_pre_norm = expected_pre_norm + float(index + 1)
+    assert torch.equal(out.hidden_states, expected_pre_norm)
+    assert not torch.equal(out.hidden_states, model.model.norm(expected_pre_norm))
+
+
 def test_generate_batch_drops_hf_flags_for_custom_backbone():
     # The custom backbone raises if any HF-only flag leaks through.
     target = HFDFlashTargetModel(_FakeCustomCausalLM(), target_layer_ids=[0, 2])
diff --git a/tests/unit_tests/speculative/test_dflash_tp.py b/tests/unit_tests/speculative/test_dflash_tp.py
index bbd4f35345..7e25f4dd80 100644
--- a/tests/unit_tests/speculative/test_dflash_tp.py
+++ b/tests/unit_tests/speculative/test_dflash_tp.py
@@ -22,15 +22,23 @@
 validated on the server.
 """
 
+from datetime import timedelta
 from types import SimpleNamespace
 
 import pytest
 import torch
+import torch.multiprocessing as mp
 import torch.nn as nn
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
 
 import nemo_automodel.recipes.llm.train_dflash as train_dflash
-from nemo_automodel.components.speculative.dflash.core import DFlashStepMetrics, DFlashTrainerModule, _to_full_tensor
+from nemo_automodel.components.loss.dllm_loss import DFlashDecayLoss
+from nemo_automodel.components.speculative.dflash.core import (
+    DFlashStepMetrics,
+    DFlashTrainerModule,
+    NoValidAnchorsError,
+    _to_full_tensor,
+)
 from nemo_automodel.components.speculative.dflash.domino_core import DominoStepMetrics, DominoTrainerModule
 from nemo_automodel.components.speculative.dflash.draft_qwen3 import Qwen3DFlashDraftModel
 from nemo_automodel.components.speculative.dflash.registry import resolve_dflash_draft_spec
@@ -96,9 +104,162 @@ def _tp_target_modules():
     parallelize_module(
         embed, mesh, RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate(), use_local_output=False)
     )
+    lm_head.requires_grad_(False)
+    embed.requires_grad_(False)
     return lm_head, embed
 
 
+def _run_two_rank_materialized_head(rank: int, init_file: str) -> None:
+    """Exercise real TP weight gathering and fused draft-gradient parity."""
+    import torch.distributed as dist
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.tensor import Shard
+    from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module
+
+    dist.init_process_group("gloo", init_method=f"file://{init_file}", rank=rank, world_size=2)
+    try:
+        torch.manual_seed(23)
+        head = nn.Linear(HIDDEN, VOCAB, bias=False)
+        reference_weight = head.weight.detach().clone()
+        mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("tp",))
+        parallelize_module(head, mesh, ColwiseParallel(output_layouts=Shard(-1), use_local_output=False))
+        head.requires_grad_(False)
+        embed = nn.Embedding(VOCAB, HIDDEN).requires_grad_(False)
+        trainer = DFlashTrainerModule(
+            draft_model=nn.Identity(),
+            target_lm_head=head,
+            target_embed_tokens=embed,
+            mask_token_id=MASK_ID,
+            block_size=BLOCK_SIZE,
+            attention_backend="sdpa",
+            num_anchors=2,
+            use_fused_linear_ce=True,
+            linear_ce_chunk_size=3,
+        )
+
+        gathered_weight, gathered_bias = trainer._materialize_frozen_lm_head(torch.device("cpu"))
+        torch.testing.assert_close(gathered_weight, reference_weight)
+        assert gathered_bias is None
+
+        B, N = 2, 2
+        T = N * (BLOCK_SIZE - 1)
+        hidden = torch.randn(B, T, HIDDEN, requires_grad=True)
+        hidden_ref = hidden.detach().clone().requires_grad_(True)
+        target_ids = torch.randint(0, VOCAB, (B, T))
+        block_mask = torch.ones(B, T)
+        loss_fn = DFlashDecayLoss(loss_gamma=4.0, chunk_size=3, normalize="mean")
+        actual, actual_correct = loss_fn.forward_fused_with_correct(
+            hidden,
+            gathered_weight,
+            target_ids,
+            block_mask,
+            block_size=BLOCK_SIZE,
+        )
+        expected_logits = torch.nn.functional.linear(hidden_ref, reference_weight)
+        expected = loss_fn(expected_logits, target_ids, block_mask, block_size=BLOCK_SIZE)
+
+        torch.testing.assert_close(actual.total_loss, expected.total_loss)
+        assert torch.equal(actual_correct, expected_logits.argmax(dim=-1) == target_ids)
+        actual.total_loss.backward()
+        expected.total_loss.backward()
+        torch.testing.assert_close(hidden.grad, hidden_ref.grad, atol=1e-6, rtol=1e-5)
+    finally:
+        dist.destroy_process_group()
+
+
+class _AncestorOwnedFrozenHead(nn.Module):
+    """Target root whose child head parameter is owned by root-level FSDP2."""
+
+    def __init__(self):
+        super().__init__()
+        self.lm_head = nn.Linear(HIDDEN, VOCAB, bias=False)
+
+
+def _run_two_rank_fsdp_ancestor_head_forward(rank: int, init_file: str) -> None:
+    """Exercise full DFlash forwards against an ancestor-owned FSDP2 head."""
+    import torch.distributed as dist
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.fsdp import FSDPModule, fully_shard
+    from torch.distributed.tensor import DTensor
+
+    dist.init_process_group(
+        "gloo",
+        init_method=f"file://{init_file}",
+        rank=rank,
+        world_size=2,
+        timeout=timedelta(seconds=60),
+    )
+    try:
+        torch.manual_seed(31)
+        target_root = _AncestorOwnedFrozenHead().requires_grad_(False)
+        expected_head_weight = target_root.lm_head.weight.detach().clone()
+        mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("dp",))
+
+        # Only the ancestor is an FSDP unit. The child lm_head has no forward
+        # hook that may be called independently; its parameter is nevertheless
+        # a sharded DTensor owned by the ancestor's FSDP state.
+        fully_shard(target_root, mesh=mesh)
+        assert isinstance(target_root, FSDPModule)
+        assert isinstance(target_root.lm_head.weight, DTensor)
+        assert not isinstance(target_root.lm_head, FSDPModule)
+
+        embed = nn.Embedding(VOCAB, HIDDEN).requires_grad_(False)
+        trainer = DFlashTrainerModule(
+            draft_model=_draft_model(),
+            target_lm_head=target_root.lm_head,
+            target_embed_tokens=embed,
+            mask_token_id=MASK_ID,
+            block_size=BLOCK_SIZE,
+            attention_backend="sdpa",
+            num_anchors=4,
+            loss_decay_gamma=4.0,
+            use_fused_linear_ce=True,
+            linear_ce_chunk_size=2,
+        )
+
+        gathered_weight, gathered_bias = trainer._materialize_frozen_lm_head(torch.device("cpu"))
+        torch.testing.assert_close(gathered_weight, expected_head_weight)
+        assert gathered_bias is None
+
+        input_ids = torch.randint(0, VOCAB - 1, (1, 12))
+        hidden = torch.randn(1, 12, len(TARGET_LAYER_IDS) * HIDDEN)
+        loss_mask = torch.ones(1, 12) if rank == 0 else torch.zeros(1, 12)
+        # Rank 0 constructs four blocks (12 predicted positions, six CE
+        # chunks). Rank 1's only legal candidate is the final possible anchor,
+        # whose three continuation positions are supervised, so it constructs
+        # one block (three predicted positions, two chunks).
+        if rank == 1:
+            loss_mask[:, 8:] = 1
+
+        metrics = trainer(input_ids=input_ids, hidden_states=hidden, loss_mask=loss_mask)
+        assert torch.isfinite(metrics.loss)
+        assert metrics.valid_blocks.item() == (4 if rank == 0 else 1)
+        metrics.loss.backward()
+        draft_grads = [parameter.grad for parameter in trainer.draft_model.parameters() if parameter.grad is not None]
+        assert draft_grads
+        assert all(torch.isfinite(grad).all() for grad in draft_grads)
+        assert target_root.lm_head.weight.grad is None
+
+        block_counts = torch.tensor([metrics.valid_blocks.item()], dtype=torch.int64)
+        gathered_counts = [torch.empty_like(block_counts) for _ in range(2)]
+        dist.all_gather(gathered_counts, block_counts)
+        assert [int(count.item()) for count in gathered_counts] == [4, 1]
+
+        # A data-dependent empty-anchor rank must still participate in the
+        # ancestor-owned head's full_tensor collective before raising. Rank 0
+        # completes its unequal local work, then both ranks reach the barrier.
+        no_anchor_mask = torch.ones(1, 12) if rank == 0 else torch.zeros(1, 12)
+        if rank == 0:
+            result = trainer(input_ids=input_ids, hidden_states=hidden, loss_mask=no_anchor_mask)
+            assert torch.isfinite(result.loss)
+        else:
+            with pytest.raises(NoValidAnchorsError):
+                trainer(input_ids=input_ids, hidden_states=hidden, loss_mask=no_anchor_mask)
+        dist.barrier()
+    finally:
+        dist.destroy_process_group()
+
+
 # --------------------------------------------------------------------------- #
 # _to_full_tensor
 # --------------------------------------------------------------------------- #
@@ -121,10 +282,45 @@ def test_to_full_tensor_gathers_vocab_sharded_dtensor(single_rank_pg):
     torch.testing.assert_close(out, full)
 
 
+def test_materialized_frozen_head_moves_to_compute_device():
+    head = nn.Linear(HIDDEN, VOCAB, bias=True).requires_grad_(False)
+    trainer = DFlashTrainerModule(
+        draft_model=nn.Identity(),
+        target_lm_head=head,
+        target_embed_tokens=nn.Embedding(VOCAB, HIDDEN).requires_grad_(False),
+        mask_token_id=MASK_ID,
+        block_size=BLOCK_SIZE,
+        use_fused_linear_ce=True,
+    )
+
+    weight, bias = trainer._materialize_frozen_lm_head(torch.device("meta"))
+
+    assert weight.device.type == "meta"
+    assert bias is not None and bias.device.type == "meta"
+
+
+def test_materialized_frozen_head_treats_empty_bias_placeholder_as_no_bias():
+    head = nn.Linear(HIDDEN, VOCAB, bias=False).requires_grad_(False)
+    head.bias = nn.Parameter(torch.empty(0), requires_grad=False)
+    trainer = DFlashTrainerModule(
+        draft_model=nn.Identity(),
+        target_lm_head=head,
+        target_embed_tokens=nn.Embedding(VOCAB, HIDDEN).requires_grad_(False),
+        mask_token_id=MASK_ID,
+        block_size=BLOCK_SIZE,
+        use_fused_linear_ce=True,
+    )
+
+    _, bias = trainer._materialize_frozen_lm_head(torch.device("cpu"))
+
+    assert bias is None
+
+
 # --------------------------------------------------------------------------- #
 # Trainer with a tensor-parallel target lm_head + embed_tokens
 # --------------------------------------------------------------------------- #
-def test_trainer_runs_with_tensor_parallel_target(single_rank_pg):
+@pytest.mark.parametrize("use_fused_linear_ce", [False, True])
+def test_trainer_runs_with_tensor_parallel_target(single_rank_pg, use_fused_linear_ce):
     """A column-parallel lm_head and vocab-parallel embed_tokens return DTensors;
     the trainer must gather them, keep them non-registered (so DDP sees only the
     draft), run a finite forward, and flow gradients to the draft."""
@@ -141,6 +337,8 @@ def test_trainer_runs_with_tensor_parallel_target(single_rank_pg):
         attention_backend="sdpa",
         num_anchors=8,
         loss_decay_gamma=7.0,
+        use_fused_linear_ce=use_fused_linear_ce,
+        linear_ce_chunk_size=3,
     )
 
     # The frozen target modules are non-registered: a DDP-wrapped trainer must
@@ -164,6 +362,21 @@ def test_trainer_runs_with_tensor_parallel_target(single_rank_pg):
     assert grad > 0
 
 
+def test_materialized_tp_head_two_rank_loss_and_gradient_parity(tmp_path):
+    """Two real ranks gather the frozen TP weight once and match the dense reference."""
+    mp.spawn(_run_two_rank_materialized_head, args=(str(tmp_path / "dflash_tp_head"),), nprocs=2, join=True)
+
+
+def test_fsdp2_ancestor_owned_head_full_forward_with_unequal_rank_work(tmp_path):
+    """Full forwards gather a root-FSDP2-owned head before rank-local chunking."""
+    mp.spawn(
+        _run_two_rank_fsdp_ancestor_head_forward,
+        args=(str(tmp_path / "dflash_fsdp_ancestor_head"),),
+        nprocs=2,
+        join=True,
+    )
+
+
 @pytest.mark.parametrize("shift_label", [True, False])
 def test_domino_trainer_runs_with_tensor_parallel_target(single_rank_pg, shift_label):
     """Regression: Domino used to consume the TP target's ``lm_head`` /