feat(dflash): add fused CE and total anchor budget

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) |

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)

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
 

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,14 +999,57 @@ 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,
             block_size,
             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):