diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py
index 10af2366f9e..29625191ef3 100644
--- a/modelopt/torch/speculative/config.py
+++ b/modelopt/torch/speculative/config.py
@@ -167,7 +167,26 @@ class DFlashConfig(ModeloptBaseConfig):
 
     dflash_use_torch_compile: bool = ModeloptField(
         default=True,
-        description="Whether to use torch.compile on DFlash forward/loss methods.",
+        description=(
+            "Whether to torch.compile the compute-heavy parts of DFlash training: the draft "
+            "decoder stack and the DSpark TVD chunk. Costs a one-time compile on the first "
+            "training step; the draft's shapes are static, so it does not recur."
+        ),
+    )
+
+    dflash_use_flex_attention: bool = ModeloptField(
+        default=False,
+        description=(
+            "Compute the draft's attention with torch FlexAttention over a block-sparse "
+            "BlockMask instead of handing SDPA a materialized [B, 1, Q, KV] float mask. "
+            "The mask is only ~33% dense (each query block sees a context PREFIX up to its "
+            "anchor, plus its own block on the diagonal), and the dense form additionally "
+            "forces SDPA onto the cutlass memory-efficient backend -- whose only kernel is "
+            "sm80 -- because the fused backends reject arbitrary masks and cap head_dim at "
+            "256. Measured 4.0x on the draft attention (225 ms -> 56 ms fwd+bwd) at the "
+            "Gemma-4-E4B DSpark shape on B300, within bf16 rounding of the dense path. "
+            "Requires torch >= 2.5. Off by default."
+        ),
     )
 
     dflash_swa_window_size: int | None = ModeloptField(
diff --git a/modelopt/torch/speculative/dflash/dflash_model.py b/modelopt/torch/speculative/dflash/dflash_model.py
index 24f2143bf84..7d90329d4bb 100644
--- a/modelopt/torch/speculative/dflash/dflash_model.py
+++ b/modelopt/torch/speculative/dflash/dflash_model.py
@@ -49,4 +49,5 @@ def modify(self, config):
         self.dflash_report_acc = config.dflash_report_acc
         self.dflash_use_torch_compile = config.dflash_use_torch_compile
         self.dflash_swa_window_size = config.dflash_swa_window_size
+        self.dflash_use_flex_attention = config.dflash_use_flex_attention
         self.dflash_export_rope_scaling = config.dflash_export_rope_scaling
diff --git a/modelopt/torch/speculative/plugins/dflash_flex_attention.py b/modelopt/torch/speculative/plugins/dflash_flex_attention.py
new file mode 100644
index 00000000000..8ea8b12a5a6
--- /dev/null
+++ b/modelopt/torch/speculative/plugins/dflash_flex_attention.py
@@ -0,0 +1,192 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Block-sparse FlexAttention path for the DFlash/DSpark draft.
+
+The draft's attention mask is dense in shape but sparse in content. Query block ``b``
+with anchor ``a_b`` attends to
+
+    context:  kv < a_b                     -- a prefix, and anchors are sorted, so
+                                              across blocks this is a staircase
+    draft:    kv in [S + Bb, S + Bb + B)   -- the block diagonal, ``block_size`` wide
+
+Handing SDPA a materialized ``[B, 1, Q, KV]`` float mask makes it compute all of it,
+and disqualifies every fused backend on the way: PyTorch's FlashAttention kernels
+reject arbitrary masks *and* cap ``head_dim`` at 256, while a Gemma-4 draft's
+``global_head_dim`` is 512. What is left is the cutlass memory-efficient backend,
+whose only kernel is **sm80** -- an Ampere kernel on Blackwell. Profiling the
+Gemma-4-E4B DSpark run on B300 put that single backward kernel at 59% of the whole
+training step, running at roughly 2.5% of the GPU's bf16 peak.
+
+FlexAttention instead takes the mask as a predicate, compiles it into the kernel, and
+skips fully-masked tiles. Measured on B300 at the production shape
+(q[4,16,4096,512], kv[4,1->16,8192,512], 512 anchors x block 8, 33% mask density):
+
+    dense-mask SDPA   fwd 24.80 ms   fwd+bwd 225.37 ms
+    flex (this file)  fwd  6.84 ms   fwd+bwd  56.40 ms      4.0x
+
+Both are within bf16 rounding of each other on out/dq/dk/dv, including the
+fully-masked rows that invalid blocks produce.
+
+Two non-obvious requirements, both established by measurement rather than docs:
+
+* ``head_dim`` 512 overflows shared memory at FlexAttention's default tiles (263 KB
+  required against a 232 KB limit), so the tiles are pinned above ``head_dim`` 256.
+  Of the shapes that fit, only 32x32 runs at all on torch 2.11 / sm103: 64x32 and
+  64x64 fault with "misaligned address" and 128x32 with "unspecified launch
+  failure". At ``head_dim`` <= 256 the library defaults are far better than anything
+  pinned (9.8x over SDPA at 256), so they are left alone.
+* ``enable_gqa=True`` must NOT be used. Its forward matches the pre-repeated path
+  exactly, but its backward takes 568 ms -- 10x slower, and 2.5x worse than the SDPA
+  baseline it is meant to replace. K/V are repeated to the query head count first,
+  which is what HF's sdpa path does anyway.
+"""
+
+import torch
+
+__all__ = ["build_draft_block_mask", "flex_attention_forward", "is_block_mask"]
+
+# head_dim > 256 cannot use FlexAttention's default tiles (SMEM overflow); 32x32 with
+# two pipeline stages is the only combination measured to both fit and run.
+_LARGE_HEAD_DIM_KERNEL_OPTIONS = {
+    "BLOCK_M": 32,
+    "BLOCK_N": 32,
+    "BLOCK_M1": 32,
+    "BLOCK_N1": 32,
+    "BLOCK_M2": 32,
+    "BLOCK_N2": 32,
+    "num_stages": 2,
+    "num_warps": 4,
+}
+_MAX_DEFAULT_TILE_HEAD_DIM = 256
+
+# Block granularity of the BlockMask itself. Finer granularity tracks the true 33%
+# density more closely (64 -> 65.7% sparsity vs 128 -> 64.2%) and measured 58.0 ms
+# against 61.4 ms; 32 is marginally better again but doubles the metadata for ~1 ms.
+_PINNED_TILE_MASK_BLOCK_SIZE = 64
+_DEFAULT_TILE_MASK_BLOCK_SIZE = 128
+
+
+def _mask_block_size(head_dim):
+    """Mask granularity, which is coupled to the kernel tiles and cannot be chosen freely.
+
+    FlexAttention requires the BlockMask's block size to be divisible by the kernel's
+    BLOCK_M/BLOCK_N, and raises "Q and KV block size must be divisible by BLOCK_M and
+    BLOCK_N" otherwise. So:
+
+    * head_dim > 256 pins 32x32 tiles (SMEM), so 64 is safe -- and finer granularity
+      tracks the true mask density better (65.7% sparsity vs 64.2% at 128, measured
+      58.0 ms vs 61.4 ms).
+    * head_dim <= 256 leaves the tiles to Inductor's autotuner, which may pick up to
+      128. Only FlexAttention's own default of 128 is guaranteed compatible with every
+      choice it can make.
+    """
+    if head_dim > _MAX_DEFAULT_TILE_HEAD_DIM:
+        return _PINNED_TILE_MASK_BLOCK_SIZE
+    return _DEFAULT_TILE_MASK_BLOCK_SIZE
+
+_flex_attention_compiled = None
+_create_block_mask_compiled = None
+
+
+def _flex_ops():
+    """Resolve and compile the FlexAttention entry points once per process."""
+    global _flex_attention_compiled, _create_block_mask_compiled
+    if _flex_attention_compiled is None:
+        from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+
+        # dynamic=False: shapes are fixed by (batch, num_anchors, block_size, seq_len)
+        # within a run, and dynamic shapes measurably deoptimize the generated kernel.
+        _flex_attention_compiled = torch.compile(flex_attention, dynamic=False)
+        _create_block_mask_compiled = torch.compile(create_block_mask, dynamic=False)
+    return _flex_attention_compiled, _create_block_mask_compiled
+
+
+def is_block_mask(mask) -> bool:
+    """True if ``mask`` is a FlexAttention ``BlockMask`` rather than a dense tensor."""
+    if mask is None or torch.is_tensor(mask):
+        return False
+    try:
+        from torch.nn.attention.flex_attention import BlockMask
+    except ImportError:
+        return False
+    return isinstance(mask, BlockMask)
+
+
+def build_draft_block_mask(
+    seq_len, anchor_positions, block_keep_mask, n_blocks, block_size, window, device,
+    head_dim,
+):
+    """BlockMask equivalent of ``HFDFlashModel._build_draft_attention_mask``.
+
+    Same predicate, expressed for FlexAttention instead of materialized. Rebuilt every
+    step because anchors are resampled on every forward; the compiled builder costs
+    ~0.1 ms, against ~1.5 ms to materialize the dense mask it replaces.
+    """
+    _, create_block_mask = _flex_ops()
+    bsz = anchor_positions.shape[0]
+    q_len = n_blocks * block_size
+    kv_len = seq_len + q_len
+    # Indexed inside mask_mod, which runs under vmap -- keep them on-device and integral.
+    anchors = anchor_positions.to(device=device, dtype=torch.int32)
+    keep = block_keep_mask.to(device=device, dtype=torch.bool)
+
+    def mask_mod(b, h, q_idx, kv_idx):
+        q_block = q_idx // block_size
+        anchor = anchors[b, q_block]
+        is_ctx = kv_idx < seq_len
+        ctx_ok = is_ctx & (kv_idx < anchor)
+        if window is not None:
+            # Same sliding window as the dense path: measured against the query's REAL
+            # position (anchor + position-in-block), not its index in the draft block.
+            ctx_ok = ctx_ok & (kv_idx > anchor + (q_idx % block_size) - window)
+        draft_ok = (~is_ctx) & (q_block == (kv_idx - seq_len) // block_size)
+        return (ctx_ok | draft_ok) & keep[b, q_block]
+
+    return create_block_mask(
+        mask_mod, bsz, None, q_len, kv_len, device=device,
+        BLOCK_SIZE=_mask_block_size(head_dim),
+    )
+
+
+def _repeat_kv(x, n_rep):
+    """HF's ``repeat_kv``: [B, n_kv, S, D] -> [B, n_kv * n_rep, S, D]."""
+    if n_rep == 1:
+        return x
+    b, h, s, d = x.shape
+    return x[:, :, None].expand(b, h, n_rep, s, d).reshape(b, h * n_rep, s, d)
+
+
+def flex_attention_forward(query, key, value, block_mask, scaling):
+    """FlexAttention with the draft's BlockMask. Returns ``[B, q_len, n_heads, head_dim]``.
+
+    The layout matches what HF's ``sdpa_attention_forward`` returns so the caller's
+    ``reshape(bsz, q_len, -1)`` is unchanged.
+    """
+    flex_attention, _ = _flex_ops()
+    head_dim = query.shape[-1]
+    n_rep = query.shape[1] // key.shape[1]
+    kernel_options = (
+        _LARGE_HEAD_DIM_KERNEL_OPTIONS if head_dim > _MAX_DEFAULT_TILE_HEAD_DIM else None
+    )
+    attn_output = flex_attention(
+        query,
+        _repeat_kv(key, n_rep),
+        _repeat_kv(value, n_rep),
+        block_mask=block_mask,
+        scale=scaling,
+        kernel_options=kernel_options,
+    )
+    return attn_output.transpose(1, 2).contiguous()
diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index b025b3088a2..4d9c99f72c1 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -568,6 +568,11 @@ def modify(self, config):
 
         self.is_quantized = False
         self._num_anchors = self.dflash_num_anchors
+        # Opt-in Inductor fusion of the draft stack; see DFlashModule._body for why this is
+        # only affordable now that the block count is static.
+        self.dflash_module._dflash_compile_stack = bool(
+            getattr(self, "dflash_use_torch_compile", False)
+        )
 
     def _build_draft_module(self, dflash_config):
         """Build the draft module. Subclasses override to use an augmented module."""
@@ -584,6 +589,24 @@ def _sample_anchor_positions(self, seq_len, loss_mask, device):
 
         Returns (anchor_positions [B, N], block_keep_mask [B, N]).
 
+        ``N`` is fixed by the config and the sequence length, never by the batch
+        contents. It used to be ``min(num_anchors, valid_counts.max() - 1)``, i.e. a
+        function of the longest answer in the batch. ``n_blocks`` sets the draft's
+        ``q_len`` and ``kv_len``, and both the block-mask builder and the attention are
+        ``torch.compile(..., dynamic=False)``, so every distinct value cost a fresh
+        170-300 s recompile -- and new values kept appearing indefinitely, since any
+        batch whose answers are all shorter than the cap mints one. A 3-node production
+        run spent 1360 of its first 1510 training seconds frozen in recompilation
+        against a 0.6 s/step steady state.
+
+        Which anchors get sampled is unchanged: ``cap`` below still applies the old
+        data-dependent bound, just on-device and before the sort rather than as a slice
+        width. The surplus columns carry ``keep=False``, which every consumer already
+        handles because rows shorter than the batch maximum have always produced them:
+        the losses normalize by the weight sum derived from ``block_keep_mask``, and
+        ``mask_mod`` ands the same flag in, so FlexAttention skips those tiles instead
+        of computing them.
+
         TODO: Fix the random seed per epoch (change between epochs) so that anchor
         positions are deterministic within an epoch. This would allow caching the derived
         masks and position IDs across steps while preserving the same data augmentation
@@ -596,27 +619,34 @@ def _sample_anchor_positions(self, seq_len, loss_mask, device):
 
         valid = loss_mask[:, : max_anchor + 1] > 0.5
         valid_counts = valid.sum(dim=1)
-        max_n = min(num_anchors, int(valid_counts.max().item()) - 1)
 
-        if max_n <= 0:
-            # No valid anchors — return empty
-            anchors = torch.zeros(bsz, 1, dtype=torch.long, device=device)
-            keep = torch.zeros(bsz, 1, dtype=torch.bool, device=device)
-            return anchors, keep
+        # Static. Bounded by max_anchor + 1 as well as num_anchors so that short sequences
+        # (unit tests, short-context recipes) cannot ask for more columns than exist.
+        max_n = min(num_anchors, max_anchor + 1)
+
+        # The bound the old code computed, kept on-device so the shapes above stay static
+        # and the .item() sync is gone. valid_counts <= max_anchor + 1 by construction, so
+        # clamping to max_n leaves this equal to min(num_anchors, valid_counts.max() - 1).
+        cap = (valid_counts.max() - 1).clamp(min=0, max=max_n)
 
         indices = torch.arange(max_anchor + 1, device=device).unsqueeze(0).expand(bsz, -1)
-        masked_indices = torch.where(valid, indices, torch.tensor(seq_len + 1, device=device))
+        fill = torch.tensor(seq_len + 1, device=device)
+        masked_indices = torch.where(valid, indices, fill)
 
         random_vals = torch.rand(bsz, max_anchor + 1, device=device)
         random_vals = torch.where(valid, random_vals, torch.tensor(2.0, device=device))
 
         _, sorted_idx = random_vals.sort(dim=1)
         gathered = torch.gather(masked_indices, 1, sorted_idx)
-        anchors = gathered[:, :max_n].sort(dim=1).values
 
-        keep = torch.arange(max_n, device=device).unsqueeze(0) < valid_counts.unsqueeze(1).clamp(
-            max=max_n
-        )
+        # Blank past `cap` BEFORE sorting, not after. Sorting a wider slice would pull
+        # anchors from beyond the old bound into the front of the row and silently change
+        # which positions are trained on; blanking first reproduces the old
+        # ``gathered[:, :cap].sort()`` exactly and leaves the surplus as fill.
+        cols = torch.arange(max_n, device=device).unsqueeze(0)
+        anchors = torch.where(cols < cap, gathered[:, :max_n], fill).sort(dim=1).values
+
+        keep = cols < torch.minimum(valid_counts, cap).unsqueeze(1)
         anchors = torch.where(keep, anchors, torch.tensor(0, dtype=torch.long, device=device))
         return anchors, keep
 
@@ -667,6 +697,24 @@ def _build_draft_attention_mask(
         q_len = n_blocks * block_size
         kv_len = seq_len + q_len
 
+        # FlexAttention consumes the same predicate as a BlockMask and skips the ~67% of
+        # tiles that are entirely masked; DFlashAttention routes a BlockMask to
+        # flex_attention_forward. See dflash_flex_attention.py for why this matters so
+        # much here (dense mask -> sm80 memory-efficient kernel -> 59% of the step).
+        if getattr(self, "dflash_use_flex_attention", False):
+            from .dflash_flex_attention import build_draft_block_mask
+
+            return build_draft_block_mask(
+                seq_len,
+                anchor_positions,
+                block_keep_mask,
+                n_blocks,
+                block_size,
+                window,
+                device,
+                head_dim=self.dflash_module.layers[0].self_attn.head_dim,
+            )
+
         q_indices = torch.arange(q_len, device=device).view(1, 1, -1, 1)
         kv_indices = torch.arange(kv_len, device=device).view(1, 1, 1, -1)
         q_block_ids = q_indices // block_size
diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py
index 4183a63f1d1..9aea6037fe0 100644
--- a/modelopt/torch/speculative/plugins/hf_dspark.py
+++ b/modelopt/torch/speculative/plugins/hf_dspark.py
@@ -78,7 +78,44 @@
 __all__ = ["HFDSparkModel"]
 
 
-def _tvd_per_token(final_logits, teacher_logits, chunk_size=1024):
+def _tvd_chunk(a, b):
+    """Per-token TVD for one row chunk: ``(softmax(a) - softmax(b)).abs().sum(-1)``."""
+    return (
+        (torch.softmax(a.float(), dim=-1) - torch.softmax(b.float(), dim=-1)).abs().sum(dim=-1)
+    )
+
+
+# torch.compile of _tvd_chunk, built once per process and reused. Eager, this chain is
+# six separate passes over [chunk, vocab] tensors (two bf16->fp32 casts, two softmaxes, a
+# subtract, an abs) that are 1 GB each at the Gemma-4 shape; Inductor fuses it into a
+# couple of kernels. Profiling put softmax + the surrounding elementwise ops at ~50% of
+# the training step once FlexAttention removed the attention bottleneck.
+_compiled_tvd_chunk = None
+_tvd_compile_failed = False
+
+
+def _get_tvd_chunk(use_compile: bool):
+    """Return the TVD chunk fn, compiled when asked for and when compilation succeeds.
+
+    Deliberately NOT wrapped in ``torch._dynamo.config.suppress_errors = True`` (which the
+    Eagle plugin sets globally): that turns a compile failure into a silent fallback to
+    eager, i.e. a performance feature that reports success while doing nothing. Here a
+    failure is warned about once and then remembered.
+    """
+    global _compiled_tvd_chunk, _tvd_compile_failed
+    if not use_compile or _tvd_compile_failed:
+        return _tvd_chunk
+    if _compiled_tvd_chunk is None:
+        try:
+            _compiled_tvd_chunk = torch.compile(_tvd_chunk, dynamic=False, fullgraph=True)
+        except Exception as exc:
+            _tvd_compile_failed = True
+            logger.warning("torch.compile of the DSpark TVD chunk failed (%s); using eager.", exc)
+            return _tvd_chunk
+    return _compiled_tvd_chunk
+
+
+def _tvd_per_token(final_logits, teacher_logits, chunk_size=1024, chunk_fn=None):
     """Total-variation distance ||softmax(a)-softmax(b)||_1 / ... per token, memory-lean.
 
     Materializing both [N, vocab] float32 softmax tensors at once OOMs at large
@@ -86,16 +123,36 @@ def _tvd_per_token(final_logits, teacher_logits, chunk_size=1024):
     gradient-checkpoint each chunk so the wide softmaxes are recomputed in backward
     rather than held — peak memory ~ chunk_size*vocab instead of N*vocab. The math
     is identical to ``(softmax(final)-softmax(teacher)).abs().sum(-1)``.
-    """
-
-    def _chunk(a, b):
-        return (
-            (torch.softmax(a.float(), dim=-1) - torch.softmax(b.float(), dim=-1)).abs().sum(dim=-1)
-        )
 
+    Chunking goes through ``Tensor.split``, NOT ``final_logits[i : i + chunk_size]``.
+    Both produce the same views over the same rows, so the forward values are
+    identical — but the backward graphs are not, and the difference is large. A slice
+    per chunk creates one ``SliceBackward0`` each, and every one of those allocates a
+    zero tensor of the FULL [N, vocab] shape and scatters its own chunk's gradient
+    into it, so the cost is O(n_chunks * N * vocab). ``split`` creates a single
+    ``SplitBackward0`` whose backward is one ``cat``, i.e. O(N * vocab).
+
+    At the Gemma-4-E4B shape (N = bsz 4 * n_blocks 512 * block_size 8 = 16384,
+    vocab = 262144 -> 8 GiB per [N, vocab] bf16 tensor, 16 chunks) that is not a
+    micro-optimization: measured on a B300, backward went 198.0 ms -> 107.9 ms, and
+    kernel attribution on the training profile put 93.2 ms/step -- 15% of a 614 ms
+    step, the second-largest item after the DDP all-reduce -- on ``SliceBackward0``.
+    Peak memory is unchanged (41.9 -> 42.1 GiB). Outputs and gradients are bitwise
+    identical: verified elementwise at the production shape for the returned
+    per-token TVD, grad(hidden) and grad(markov_w2.weight).
+
+    Do NOT substitute ``torch.chunk`` or ``torch.tensor_split``. Those split into a
+    fixed NUMBER of pieces and so pick different boundaries (ceil(N/k)), which changes
+    the shapes handed to ``chunk_fn`` -- and ``chunk_fn`` may be a
+    ``torch.compile(..., dynamic=False)`` build that recompiles per shape.
+    """
+    _chunk = chunk_fn or _tvd_chunk
     outs = []
-    for i in range(0, final_logits.size(0), chunk_size):
-        a, b = final_logits[i : i + chunk_size], teacher_logits[i : i + chunk_size]
+    for a, b in zip(
+        final_logits.split(chunk_size, dim=0),
+        teacher_logits.split(chunk_size, dim=0),
+        strict=True,
+    ):
         if torch.is_grad_enabled() and a.requires_grad:
             outs.append(torch.utils.checkpoint.checkpoint(_chunk, a, b, use_reentrant=False))
         else:
@@ -141,9 +198,17 @@ def get_exporter(self):
 
         return DSparkExporter(self)
 
-    def _apply_markov_head(self, hidden, backbone_logits, input_ids, anchor_positions, n_blocks):
+    def _apply_markov_head(
+        self, hidden, backbone_logits, input_ids, anchor_positions, n_blocks, inplace=False
+    ):
         """Add the Markov transition bias to the backbone base logits.
 
+        ``inplace`` folds the bias into ``backbone_logits`` instead of allocating a third
+        [B, N, bs, vocab] tensor (8.6 GB at the Gemma-4 shape). Safe for autograd -- neither
+        ``lm_head`` nor ``markov_w2`` saves its OUTPUT for backward -- but it leaves
+        ``backbone_logits`` holding the corrected logits, so the caller must not still need
+        the uncorrected ones (they are only used for the ``base_accuracy`` metric).
+
         Returns ``(final_logits [B, N, bs, V], confidence_logits [B, N, bs] | None)``.
         """
         bsz, seq_len = input_ids.shape
@@ -160,7 +225,7 @@ def _apply_markov_head(self, hidden, backbone_logits, input_ids, anchor_position
         prev_ids = torch.gather(input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, prev_idx)
 
         bias = self.dflash_module.compute_markov_bias(prev_ids, hidden4d)
-        final4d = base4d + bias
+        final4d = base4d.add_(bias) if inplace else base4d + bias
 
         confidence_logits = None
         if self.dflash_module.use_confidence_head:
@@ -176,13 +241,24 @@ def _compute_dspark_loss(
         anchor_positions,
         block_keep_mask,
         loss_mask,
-        target_model_logits,
+        target_model_logits=None,
+        teacher_hidden=None,
     ):
         """Compute the three-term DSpark loss (CE + TVD + confidence BCE) and metrics.
 
         Uses next-token (shift_label) alignment: block position k predicts the token
         at anchor+k+1; the aligned target distribution is the base model's own
         next-token distribution at position anchor+k (= label index - 1).
+
+        The teacher distribution can arrive two ways. ``teacher_hidden`` ([B, seq, H], the
+        base model's post-final-norm hidden) is preferred: only the N*block_size rows the
+        loss actually reads are gathered and projected, so the full-sequence
+        [B, seq, vocab] tensor is never built and never gathered out of -- at the Gemma-4
+        shape that is ~17 GB/step of memory traffic. ``target_model_logits`` is the
+        fallback for producers that hand over logits directly.
+
+        ``backbone_logits`` may be None, which skips the ``base_accuracy`` diagnostic (and
+        with it a second full-vocab argmax); see ``dflash_report_acc``.
         """
         bsz, seq_len = input_ids.shape
         bs = self.dflash_block_size
@@ -217,7 +293,7 @@ def _compute_dspark_loss(
             weight_mask = weight_mask * decay
 
         flat_final = final_logits.reshape(-1, vocab)
-        flat_base = backbone_logits.reshape(-1, vocab)
+        flat_base = None if backbone_logits is None else backbone_logits.reshape(-1, vocab)
         flat_targets = target_ids.reshape(-1)
         flat_weights = weight_mask.reshape(-1)
         valid_count = flat_weights.sum() + 1e-6
@@ -225,15 +301,37 @@ def _compute_dspark_loss(
         # Aligned target distribution: base-model logits that predict token anchor+k+1
         # sit at position anchor+k (= label index - 1).
         teacher_indices = (safe_label_indices - 1).clamp(min=0)
-        teacher_logits = torch.gather(
-            target_model_logits.unsqueeze(1).expand(-1, n_blocks, -1, -1),
-            2,
-            teacher_indices.unsqueeze(-1).expand(-1, -1, -1, vocab),
-        )
-        flat_teacher = teacher_logits.reshape(-1, vocab).detach()
+        with torch.no_grad():
+            if teacher_hidden is not None:
+                hdim = teacher_hidden.size(-1)
+                gathered_hidden = torch.gather(
+                    teacher_hidden.unsqueeze(1).expand(-1, n_blocks, -1, -1),
+                    2,
+                    teacher_indices.unsqueeze(-1).expand(-1, -1, -1, hdim),
+                )
+                flat_teacher = self._base_model_lm_head(gathered_hidden.reshape(-1, hdim))
+            else:
+                if target_model_logits is None:
+                    raise ValueError(
+                        "DSpark loss needs the base distribution: pass teacher_hidden "
+                        "(preferred) or target_model_logits."
+                    )
+                flat_teacher = torch.gather(
+                    target_model_logits.unsqueeze(1).expand(-1, n_blocks, -1, -1),
+                    2,
+                    teacher_indices.unsqueeze(-1).expand(-1, -1, -1, vocab),
+                ).reshape(-1, vocab)
+        flat_teacher = flat_teacher.detach()
 
         if valid_count <= 1.0:
-            loss = flat_final.sum() * 0.0
+            # Every draft parameter must receive a gradient, not just the ones behind
+            # final_logits: the confidence head hangs off compute_confidence_logits, which
+            # this branch skips, so `flat_final.sum() * 0` alone leaves it unused and DDP
+            # with find_unused_parameters=False aborts the run on the first such batch.
+            # Mirrors the same guard in forward()'s no-valid-anchor early return.
+            loss = flat_final.sum() * 0.0 + sum(
+                p.sum() for p in self.dflash_module.parameters()
+            ) * 0.0
             metrics = {"ce_loss": 0.0, "l1_loss": 0.0, "confidence_loss": 0.0, "base_accuracy": 0.0}
             return loss, 0.0, metrics
 
@@ -243,7 +341,11 @@ def _compute_dspark_loss(
 
         # Term 2: total-variation distance between the corrected draft and target.
         # Chunked + checkpointed to avoid materializing two [N, vocab] softmaxes at once.
-        l1_per_token = _tvd_per_token(flat_final, flat_teacher)
+        l1_per_token = _tvd_per_token(
+            flat_final,
+            flat_teacher,
+            chunk_fn=_get_tvd_chunk(getattr(self, "dflash_use_torch_compile", False)),
+        )
         l1_loss = (l1_per_token * flat_weights).sum() / valid_count
 
         # Term 3: confidence head BCE against the analytical accept rate c* = 1 - 0.5*TVD.
@@ -264,19 +366,25 @@ def _compute_dspark_loss(
         with torch.no_grad():
             eval_count = binary_eval_mask.sum() + 1e-6
             keep = binary_eval_mask > 0.5
-            accuracy = (
-                ((flat_final.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count
-            ).item()
-            base_accuracy = (
-                ((flat_base.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count
-            ).item()
+            acc = ((flat_final.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count
+            base_acc = (
+                acc.new_zeros(())
+                if flat_base is None
+                else ((flat_base.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count
+            )
+            # ONE device sync for all five scalars. Each .item() is a full synchronize, and
+            # five of them per step chop up the window in which DDP's all-reduce can hide
+            # behind backward -- measured comm exposure is 18% of the step post-FlexAttention.
+            acc_v, base_acc_v, ce_v, l1_v, conf_v = torch.stack(
+                [acc, base_acc, ce_loss.detach(), l1_loss.detach(), confidence_loss.detach()]
+            ).tolist()
             metrics = {
-                "ce_loss": ce_loss.detach().item(),
-                "l1_loss": l1_loss.detach().item(),
-                "confidence_loss": float(confidence_loss.detach().item()),
-                "base_accuracy": base_accuracy,
+                "ce_loss": ce_v,
+                "l1_loss": l1_v,
+                "confidence_loss": conf_v,
+                "base_accuracy": base_acc_v,
             }
-        return loss, accuracy, metrics
+        return loss, acc_v, metrics
 
     def forward(
         self,
@@ -336,9 +444,14 @@ def forward(
                 self._base_model_norm,
                 self._base_model_lm_head,
                 need_logits=True,
+                # Hand back the normed hidden instead of full-sequence logits; the loss
+                # projects only the rows it reads. Producers that supply base_model_logits
+                # directly still come back with logits and take the fallback path.
+                defer_lm_head=True,
             )
             target_hidden = base_outputs.target_hidden
             target_model_logits = base_outputs.logits
+            teacher_hidden = base_outputs.base_hidden
         else:
             # Call the inner base model directly (NOT super().forward(), which during
             # training runs the full DFlash pipeline). Compute target-model logits via
@@ -349,7 +462,9 @@ def forward(
                     attention_mask=attention_mask,
                     output_hidden_states=True,
                 )
-                target_model_logits = self._base_model_lm_head(base_out.last_hidden_state)
+                # lm_head is applied per-row inside the loss, not over the whole sequence.
+                teacher_hidden = base_out.last_hidden_state
+                target_model_logits = None
             offset = 1
             selected = [base_out.hidden_states[lid + offset] for lid in self.target_layer_ids]
             target_hidden = torch.cat(selected, dim=-1)  # [B, seq, num_layers * H]
@@ -399,19 +514,30 @@ def forward(
         )
 
         # 6. Backbone logits → Markov correction → three-term loss.
+        # dflash_report_acc gates the base_accuracy diagnostic (the draft's accuracy BEFORE
+        # the Markov correction). With it off, the uncorrected logits are dead after the
+        # bias is added, so the bias folds in place and a second full-vocab argmax is
+        # skipped -- two [N, vocab] tensors' worth of traffic per step.
+        report_base_acc = getattr(self, "dflash_report_acc", True)
         backbone_logits = self._base_model_lm_head(hidden).reshape(bsz, n_blocks, block_size, -1)
         final_logits, confidence_logits = self._apply_markov_head(
-            hidden, backbone_logits, input_ids, anchor_positions, n_blocks
+            hidden,
+            backbone_logits,
+            input_ids,
+            anchor_positions,
+            n_blocks,
+            inplace=not report_base_acc,
         )
         loss, accuracy, metrics = self._compute_dspark_loss(
-            backbone_logits,
+            backbone_logits if report_base_acc else None,
             final_logits,
             confidence_logits,
             input_ids,
             anchor_positions,
             block_keep_mask,
             loss_mask,
-            target_model_logits,
+            target_model_logits=target_model_logits,
+            teacher_hidden=teacher_hidden,
         )
 
         return ModelOutput(loss=loss, logits=None, train_acc=[[accuracy]], dspark_metrics=metrics)
diff --git a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py
index 0f9713948fd..1a81dbdf519 100644
--- a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py
+++ b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py
@@ -536,14 +536,22 @@ def _fetch(self, sample: dict) -> EagleFetchPayload | None:
             time.sleep(0.0002)
         agent.release_xfer_handle(h)
         hidden_states = view.clone()  # copy out before /done so the gen check brackets the read
-        # /done frees the slot + reports valid; valid=False -> ring lapped us mid-read, bytes
-        # stale -> resample. A failed /done can't prove staleness, so default valid=True.
+        # /done frees the slot and reports whether the ring lapped us mid-read (stale bytes).
+        # It fails CLOSED: a /done that errors cannot prove the slot is still ours either, and
+        # the two mistakes are not symmetric. Discarding a good sample costs one resample;
+        # keeping a lapped one trains the draft on another prompt's activations, and nothing
+        # downstream catches that -- the token_ids check below compares the server's
+        # per-request record, not the bytes just read, so it passes on a mis-slotted read.
+        # A rising resample rate is at least visible in the logs; silent corruption is not.
         try:
             valid = self._http_rdma.get(
                 f"http://{host}:{port}/done", params={"req_id": rid}
             ).json()["valid"]
-        except Exception:
-            valid = True
+        except Exception as exc:
+            # Deliberately distinct from the lap message below: a wave of these is a sidecar
+            # problem, not ring pressure, and the two call for different fixes.
+            warn_rank_0(f"[streaming] /done failed for {sample['cid']} ({exc!r}); resampling")
+            return None
         if not valid:
             warn_rank_0(f"[streaming] slot lapped mid-read for {sample['cid']}; resampling")
             return None
diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py
index 9b6c459c1a2..cc1cc7b37f5 100644
--- a/modelopt/torch/speculative/plugins/modeling_dflash.py
+++ b/modelopt/torch/speculative/plugins/modeling_dflash.py
@@ -65,10 +65,19 @@ class DFlashBaseModelOutput:
 
     target_hidden: torch.Tensor  # concatenated hidden states from target layers [B, seq, N*H]
     logits: torch.Tensor | None = None  # base model logits [B, seq, vocab]
+    # Post-final-norm base hidden [B, seq, H], i.e. lm_head's input. Consumers that only
+    # need the base distribution at a handful of positions project THIS at those rows
+    # instead of materialising (and then gathering out of) full-sequence logits.
+    base_hidden: torch.Tensor | None = None
 
     @classmethod
     def from_offline_dict(
-        cls, d: dict, base_model_norm=None, base_model_lm_head=None, need_logits=False
+        cls,
+        d: dict,
+        base_model_norm=None,
+        base_model_lm_head=None,
+        need_logits=False,
+        defer_lm_head=False,
     ):
         """Construct from a dict of pre-computed base model outputs (offline training).
 
@@ -85,19 +94,25 @@ def from_offline_dict(
         to lm_head would be a corrupt distillation target).
         """
         logits = d.get("base_model_logits")
+        base_hidden = None
         if need_logits and logits is None:
+            out_hiddens = d.get("base_model_hidden_states")
+            if out_hiddens is None:
+                raise KeyError("base_model_hidden_states")
+            base_hidden = _maybe_apply_base_final_norm(out_hiddens, d, base_model_norm)
+            if defer_lm_head:
+                # Caller will project only the rows it needs; skip the full-sequence
+                # [B, seq, vocab] materialisation entirely.
+                return cls(target_hidden=d["aux_hidden_states"], base_hidden=base_hidden)
             if base_model_lm_head is None:
                 raise ValueError(
                     "need_logits=True but base_model_lm_head is None; cannot reconstruct logits."
                 )
-            out_hiddens = d.get("base_model_hidden_states")
-            if out_hiddens is None:
-                raise KeyError("base_model_hidden_states")
-            out_hiddens = _maybe_apply_base_final_norm(out_hiddens, d, base_model_norm)
-            logits = base_model_lm_head(out_hiddens)
+            logits = base_model_lm_head(base_hidden)
         return cls(
             target_hidden=d["aux_hidden_states"],
             logits=logits,
+            base_hidden=base_hidden,
         )
 
 
@@ -185,6 +200,36 @@ def _get_attn_fn(self):
         self._attn_fn = ALL_ATTENTION_FUNCTIONS.get(impl, ALL_ATTENTION_FUNCTIONS["sdpa"])
         return self._attn_fn
 
+    def _attend(self, q, k, v, attention_mask, bsz, q_len):
+        """Run attention and project, routing a FlexAttention BlockMask to the flex kernel.
+
+        ``attention_mask`` is either the dense additive [B, 1, Q, KV] tensor (HF attention
+        dispatch) or a BlockMask carrying the same predicate block-sparsely.
+        """
+        from .dflash_flex_attention import flex_attention_forward, is_block_mask
+
+        if is_block_mask(attention_mask):
+            dropout = 0.0 if not self.training else self.attention_dropout
+            if dropout:
+                raise ValueError(
+                    "FlexAttention path does not support attention_dropout > 0 "
+                    f"(got {dropout}); unset dflash_use_flex_attention."
+                )
+            attn_output = flex_attention_forward(q, k, v, attention_mask, self.scaling)
+        else:
+            attn_fn = self._get_attn_fn()
+            attn_output, _ = attn_fn(
+                self,
+                q,
+                k,
+                v,
+                attention_mask,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                scaling=self.scaling,
+                sliding_window=self.sliding_window,
+            )
+        return self.o_proj(attn_output.reshape(bsz, q_len, -1))
+
     def forward(self, hidden_states, target_hidden, position_embeddings, attention_mask=None):
         """Forward with KV injection.
 
@@ -218,20 +263,7 @@ def forward(self, hidden_states, target_hidden, position_embeddings, attention_m
         cos, sin = position_embeddings
         q, k = apply_rotary_pos_emb(q, k, cos, sin)
 
-        # Use HF's attention dispatch (handles GQA internally)
-        attn_fn = self._get_attn_fn()
-        attn_output, _ = attn_fn(
-            self,
-            q,
-            k,
-            v,
-            attention_mask,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.scaling,
-            sliding_window=self.sliding_window,
-        )
-        attn_output = attn_output.reshape(bsz, q_len, -1)
-        return self.o_proj(attn_output)
+        return self._attend(q, k, v, attention_mask, bsz, q_len)
 
 
 class DFlashGemma4Attention(DFlashAttention):
@@ -353,19 +385,7 @@ def forward(self, hidden_states, target_hidden, position_embeddings, attention_m
         cos, sin = position_embeddings
         q, k = apply_rotary_pos_emb(q, k, cos, sin)
 
-        attn_fn = self._get_attn_fn()
-        attn_output, _ = attn_fn(
-            self,
-            q,
-            k,
-            v,
-            attention_mask,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.scaling,
-            sliding_window=self.sliding_window,
-        )
-        attn_output = attn_output.reshape(bsz, q_len, -1)
-        return self.o_proj(attn_output)
+        return self._attend(q, k, v, attention_mask, bsz, q_len)
 
 
 class DFlashDecoderLayer(nn.Module):
@@ -553,9 +573,38 @@ def _init_weights(self, config):
 
     def forward(self, noise_embedding, target_hidden, position_ids, attention_mask=None):
         """Forward with feature fusion, KV injection, and position embeddings."""
+        # Lazy rotary construction mutates the module, so it stays outside the compiled
+        # region below: Dynamo would either graph-break on it or bake in the first call's
+        # state.
+        self._maybe_init_rotary_emb(device=noise_embedding.device)
+        return self._body()(noise_embedding, target_hidden, position_ids, attention_mask)
+
+    def _body(self):
+        """Return the draft stack, Inductor-compiled on first use when asked for.
+
+        The layer loop is where the step's small-kernel tail lives: at the Gemma-4-E4B
+        shape a profiled step ran ~1500 pointwise launches, 403 ``aten::copy_`` and 205
+        device-to-device memcpys, individually microseconds and collectively about a third
+        of the step. Fusing them is a compiler's job.
+
+        ``dynamic=False`` is only affordable because ``_sample_anchor_positions`` pins the
+        block count; while n_blocks tracked the batch, every new width cost a fresh
+        multi-minute compile.
+
+        Training only. Generation (AR validation, drafting) runs the same module at a
+        different and varying shape, which under ``dynamic=False`` would mint a compile per
+        length -- exactly the trade the pinned block count was introduced to avoid.
+        """
+        if not self.training or not getattr(self, "_dflash_compile_stack", False):
+            return self._forward_body
+        if getattr(self, "_compiled_body", None) is None:
+            self._compiled_body = torch.compile(self._forward_body, dynamic=False)
+        return self._compiled_body
+
+    def _forward_body(self, noise_embedding, target_hidden, position_ids, attention_mask):
+        """Feature fusion, rotary selection, the decoder stack, and the final norm."""
         hidden_states = noise_embedding
         target_hidden = self.hidden_norm(self.fc(target_hidden))
-        self._maybe_init_rotary_emb(device=hidden_states.device)
         per_kind = {
             kind: emb(hidden_states, position_ids)
             for kind, emb in getattr(self, "rotary_emb_by_kind", {}).items()
diff --git a/modelopt/torch/speculative/plugins/modeling_dspark.py b/modelopt/torch/speculative/plugins/modeling_dspark.py
index aa1a4181bbd..1dbce2768c9 100644
--- a/modelopt/torch/speculative/plugins/modeling_dspark.py
+++ b/modelopt/torch/speculative/plugins/modeling_dspark.py
@@ -142,44 +142,58 @@ def prev_token_embeddings(self, prev_ids: torch.Tensor) -> torch.Tensor:
         """Look up the Markov embedding ``W1[x_{k-1}]`` of the teacher-forced prev tokens."""
         return self.markov_w1(prev_ids.long())
 
-    def compute_markov_bias(self, prev_ids: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor:
-        """Compute the transition bias ``B_k`` added to the backbone base logits.
+    def compute_markov_latent(self, prev_ids: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor:
+        """Rank-``r`` state the transition bias is projected from, BEFORE ``markov_w2``.
+
+        Every head type ends in ``markov_w2(<width r>)``; this returns that argument so a
+        caller can defer the vocab-wide projection (which materialises a
+        ``[B, N, block_size, vocab]`` tensor) to the point where it is actually needed.
 
         Args:
             prev_ids: Teacher-forced previous-token ids per block position [B, N, block_size].
             hidden: Backbone hidden states [B, N, block_size, H] (used by gated/rnn heads).
 
         Returns:
-            Logit bias [B, N, block_size, vocab].
+            Latent [B, N, block_size, r].
         """
         prev_emb = self.prev_token_embeddings(prev_ids)  # [B, N, bs, r]
 
         if self.markov_head_type == "vanilla":
-            return self.markov_w2(prev_emb)
+            return prev_emb
 
         if self.markov_head_type == "gated":
             gate = torch.sigmoid(self.gate_proj(torch.cat([hidden, prev_emb], dim=-1)))
-            return self.markov_w2(gate.to(prev_emb.dtype) * prev_emb)
+            return gate.to(prev_emb.dtype) * prev_emb
 
         # rnn: unroll the gated recurrence over the block dimension.
         block_size = prev_ids.shape[-1]
         leading = prev_emb.shape[:-2]  # [B, N]
         state = torch.zeros(*leading, self.markov_rank, device=prev_emb.device, dtype=hidden.dtype)
-        biases = []
+        latents = []
         for k in range(block_size):
-            state, bias = self._rnn_step(state, prev_emb[..., k, :], hidden[..., k, :])
-            biases.append(bias)
-        return torch.stack(biases, dim=-2)
+            state, latent = self._rnn_step(state, prev_emb[..., k, :], hidden[..., k, :])
+            latents.append(latent)
+        return torch.stack(latents, dim=-2)
+
+    def compute_markov_bias(self, prev_ids: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor:
+        """Compute the transition bias ``B_k`` added to the backbone base logits.
+
+        Returns:
+            Logit bias [B, N, block_size, vocab].
+        """
+        return self.markov_w2(self.compute_markov_latent(prev_ids, hidden))
 
     def _rnn_step(self, state, prev_emb, hidden):
-        """One GRU-like recurrent step. Returns (new_state [.., r], bias [.., vocab])."""
+        """One GRU-like recurrent step. Returns (new_state [.., r], latent [.., r]).
+
+        The latent is ``markov_w2``'s input, not the bias -- see compute_markov_latent.
+        """
         z = torch.cat([state, prev_emb, hidden], dim=-1)
         gate_raw, candidate_raw, output_raw = self.joint_proj(z).chunk(3, dim=-1)
         gate = torch.sigmoid(gate_raw)
         candidate = torch.tanh(candidate_raw)
         new_state = gate * state + (1.0 - gate) * candidate
-        bias = self.markov_w2(torch.tanh(output_raw))
-        return new_state, bias
+        return new_state, torch.tanh(output_raw)
 
     def markov_step(self, prev_token: torch.Tensor, hidden: torch.Tensor, state=None):
         """One autoregressive Markov step (inference): bias for a single position.
@@ -204,7 +218,8 @@ def markov_step(self, prev_token: torch.Tensor, hidden: torch.Tensor, state=None
             state = torch.zeros(
                 prev_emb.shape[0], self.markov_rank, device=prev_emb.device, dtype=hidden.dtype
             )
-        state, bias = self._rnn_step(state, prev_emb, hidden)
+        state, latent = self._rnn_step(state, prev_emb, hidden)
+        bias = self.markov_w2(latent)
         return bias, state
 
     def compute_confidence_logits(
diff --git a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml
index dc849cb1ba7..899bf148c48 100644
--- a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml
+++ b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml
@@ -103,7 +103,13 @@ training:
   bf16: true
   tf32: true
   remove_unused_columns: false
-  ddp_find_unused_parameters: true
+  # false: every draft parameter gets a gradient on every batch, including the
+  # no-valid-anchor batches that skip the three loss terms -- forward() and the
+  # valid_count <= 1 branch both walk dflash_module.parameters() to keep DDP fed. True
+  # costs an extra autograd-graph traversal per iteration for nothing (torch itself warns
+  # it "did not find any unused parameters"). Precondition for this recipe:
+  # dflash_confidence_head_alpha > 0, otherwise the confidence head is never in the loss.
+  ddp_find_unused_parameters: false
   ddp_timeout: 1800
   report_to: none
 
diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py
index ca0c6f779a8..ef6831a1d96 100644
--- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py
@@ -250,3 +250,124 @@ def test_offline_forward_self_logit_distillation_recomputes_logits(self, offline
         assert hasattr(output, "logits")
         assert output.logits is not None
         assert torch.isfinite(output.loss).item()
+
+
+class TestDFlashFlexAttentionGPU:
+    """FlexAttention path must match the dense-mask SDPA path it replaces.
+
+    The block-sparse BlockMask encodes exactly the predicate
+    ``_build_draft_attention_mask`` materializes, so switching implementations may only
+    move results by float rounding -- not by a masked position becoming visible, and not
+    at the fully-masked query rows that invalid blocks produce.
+
+    These build a wider model than the rest of this file: ``get_tiny_llama`` defaults to
+    hidden_size 32 over 16 heads, i.e. head_dim 2, and FlexAttention's Triton templates
+    need head_dim >= 16 (below that Inductor finds no valid config and raises). 128/2
+    heads gives head_dim 64, the smallest standard size.
+    """
+
+    SEQ = 64
+    BASE_KWARGS = {
+        "num_hidden_layers": 4,
+        "hidden_size": 128,
+        "num_attention_heads": 2,
+        "num_key_value_heads": 1,
+        "intermediate_size": 64,
+        "max_position_embeddings": 256,
+        "vocab_size": 64,
+    }
+
+    @classmethod
+    def _model(cls, **cfg_overrides):
+        model = get_tiny_llama(**cls.BASE_KWARGS)
+        config = get_dflash_config()
+        config.update(cfg_overrides)
+        mtsp.convert(model, [("dflash", config)])
+        return model.cuda().train()
+
+    @classmethod
+    def _pair(cls, **cfg_overrides):
+        """A dense model and a flex model with identical draft weights."""
+        dense = cls._model(**cfg_overrides)
+        flex = cls._model(dflash_use_flex_attention=True, **cfg_overrides)
+        flex.dflash_module.load_state_dict(dense.dflash_module.state_dict())
+        return dense, flex
+
+    @classmethod
+    def _inputs(cls, bsz=2):
+        input_ids = torch.randint(0, cls.BASE_KWARGS["vocab_size"], (bsz, cls.SEQ), device="cuda")
+        attention_mask = torch.ones(bsz, cls.SEQ, dtype=torch.long, device="cuda")
+        return input_ids, attention_mask
+
+    def test_mask_builder_returns_block_mask(self):
+        """With the flag on, the mask builder hands back a BlockMask, not a dense tensor."""
+        pytest.importorskip("torch.nn.attention.flex_attention")
+        from modelopt.torch.speculative.plugins.dflash_flex_attention import is_block_mask
+
+        model = self._model(dflash_use_flex_attention=True)
+        mask = model._build_draft_attention_mask(
+            self.SEQ,
+            torch.tensor([[4, 8]], device="cuda"),
+            torch.tensor([[True, True]], device="cuda"),
+            2,
+            torch.float32,
+            torch.device("cuda"),
+            window=None,
+        )
+        assert is_block_mask(mask)
+        assert not torch.is_tensor(mask)
+
+    @pytest.mark.parametrize("window", [None, 8])
+    def test_matches_dense_mask_path(self, window):
+        """Loss agrees with the dense path to bf16 tolerance, with and without SWA."""
+        pytest.importorskip("torch.nn.attention.flex_attention")
+        overrides = {} if window is None else {"dflash_swa_window_size": window}
+        dense, flex = self._pair(**overrides)
+        input_ids, attention_mask = self._inputs()
+
+        # Anchors are resampled from the RNG on every forward, so both models must draw
+        # from the same seed or they would see different anchors, not different kernels.
+        torch.manual_seed(1234)
+        out_dense = dense(input_ids=input_ids, attention_mask=attention_mask)
+        torch.manual_seed(1234)
+        out_flex = flex(input_ids=input_ids, attention_mask=attention_mask)
+
+        torch.testing.assert_close(out_flex.loss, out_dense.loss, rtol=2e-2, atol=2e-2)
+
+    def test_matches_dense_mask_path_with_invalid_blocks(self):
+        """Fully-masked query rows (invalid blocks) must not diverge either.
+
+        A row of all -inf is the one place the two kernels could legitimately disagree
+        (softmax of nothing), and answer_only_loss produces such rows whenever a sample
+        has fewer valid anchors than the batch maximum.
+        """
+        pytest.importorskip("torch.nn.attention.flex_attention")
+        dense, flex = self._pair()
+        input_ids, attention_mask = self._inputs()
+        # Row 1 keeps far fewer supervised positions than row 0, so its trailing blocks
+        # come back with block_keep_mask False.
+        labels = input_ids.clone()
+        labels[1, : self.SEQ - BLOCK_SIZE] = -100
+
+        torch.manual_seed(7)
+        out_dense = dense(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+        torch.manual_seed(7)
+        out_flex = flex(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+
+        assert torch.isfinite(out_flex.loss), "flex produced a non-finite loss"
+        torch.testing.assert_close(out_flex.loss, out_dense.loss, rtol=2e-2, atol=2e-2)
+
+    def test_backward_produces_finite_grads(self):
+        """The flex path is differentiable and its grads are finite."""
+        pytest.importorskip("torch.nn.attention.flex_attention")
+        flex = self._model(dflash_use_flex_attention=True)
+        input_ids, attention_mask = self._inputs()
+
+        flex(input_ids=input_ids, attention_mask=attention_mask).loss.backward()
+        grads = [
+            p.grad
+            for p in flex.dflash_module.parameters()
+            if p.requires_grad and p.grad is not None
+        ]
+        assert grads, "no draft gradients were produced"
+        assert all(torch.isfinite(g).all() for g in grads)
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index ab5c5a57d21..336bec17e96 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -912,3 +912,175 @@ def test_multi_turn_masks_only_assistant(self, tiny_tokenizer):
         # User/system content should NOT appear in unmasked tokens
         assert "You are helpful" not in decoded
         assert "How are you?" not in decoded
+
+
+def _legacy_sample_anchor_positions(model, seq_len, loss_mask, device):
+    """Verbatim copy of the pre-static implementation, kept as the semantic reference."""
+    bs = model.dflash_block_size
+    bsz = loss_mask.shape[0]
+    max_anchor = max(seq_len - bs, 0)
+    num_anchors = getattr(model, "_num_anchors", 512)
+
+    valid = loss_mask[:, : max_anchor + 1] > 0.5
+    valid_counts = valid.sum(dim=1)
+    max_n = min(num_anchors, int(valid_counts.max().item()) - 1)
+
+    if max_n <= 0:
+        return (
+            torch.zeros(bsz, 1, dtype=torch.long, device=device),
+            torch.zeros(bsz, 1, dtype=torch.bool, device=device),
+        )
+
+    indices = torch.arange(max_anchor + 1, device=device).unsqueeze(0).expand(bsz, -1)
+    masked_indices = torch.where(valid, indices, torch.tensor(seq_len + 1, device=device))
+    random_vals = torch.rand(bsz, max_anchor + 1, device=device)
+    random_vals = torch.where(valid, random_vals, torch.tensor(2.0, device=device))
+    _, sorted_idx = random_vals.sort(dim=1)
+    gathered = torch.gather(masked_indices, 1, sorted_idx)
+    anchors = gathered[:, :max_n].sort(dim=1).values
+    keep = torch.arange(max_n, device=device).unsqueeze(0) < valid_counts.unsqueeze(1).clamp(
+        max=max_n
+    )
+    anchors = torch.where(keep, anchors, torch.tensor(0, dtype=torch.long, device=device))
+    return anchors, keep
+
+
+class TestAnchorSamplingStaticShape:
+    """n_blocks is fixed by the config, and the anchors it picks are still the legacy ones.
+
+    A data-dependent n_blocks made the draft's q_len/kv_len data-dependent, and both the
+    block-mask builder and the attention are compiled with ``dynamic=False`` -- so every
+    new batch shape cost a multi-minute recompile, indefinitely. These tests pin both
+    halves of the fix: the shape no longer moves with the data, and what the model trains
+    on did not change.
+    """
+
+    DEVICE = torch.device("cpu")
+
+    @staticmethod
+    def _model(num_anchors):
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = get_dflash_config(block_size=BLOCK_SIZE)
+        config["dflash_num_anchors"] = num_anchors
+        mtsp.convert(model, [("dflash", config)])
+        return model
+
+    @staticmethod
+    def _loss_mask(lengths, seq_len=SEQ_LEN):
+        mask = torch.zeros(len(lengths), seq_len)
+        for row, n in enumerate(lengths):
+            mask[row, :n] = 1.0
+        return mask
+
+    # Each case puts the legacy bound somewhere different: under the cap, ragged across
+    # rows, at 1, and in the two degenerate cases the old code special-cased.
+    @pytest.mark.parametrize(
+        "lengths",
+        [(13, 13), (13, 5), (9, 9), (5, 3), (2, 1), (1, 1), (0, 0), (13, 0)],
+    )
+    def test_matches_legacy_sampling(self, lengths):
+        """Same seed in, same anchors and same keep mask out -- bitwise."""
+        model = self._model(num_anchors=8)
+        loss_mask = self._loss_mask(lengths)
+
+        torch.manual_seed(1234)
+        legacy_anchors, legacy_keep = _legacy_sample_anchor_positions(
+            model, SEQ_LEN, loss_mask, self.DEVICE
+        )
+        torch.manual_seed(1234)
+        anchors, keep = model._sample_anchor_positions(SEQ_LEN, loss_mask, self.DEVICE)
+
+        n_old = legacy_keep.shape[1]
+        assert keep.shape[1] >= n_old
+        assert torch.equal(keep[:, :n_old], legacy_keep)
+        assert torch.equal(anchors[:, :n_old], legacy_anchors)
+        # Everything past the legacy bound is inert padding.
+        assert not keep[:, n_old:].any()
+        assert not anchors[:, n_old:].any()
+
+    def test_sort_is_truncated_before_it_widens(self):
+        """The surplus columns must not pull unsampled anchors to the front of the row.
+
+        Slicing to the static width *after* sorting would do exactly that: anchors past
+        the legacy bound would sort in among the kept ones and silently move which
+        positions are trained on, while every shape assertion still passed.
+        """
+        model = self._model(num_anchors=8)
+        # One long row, so the legacy bound (valid_counts.max() - 1) sits below the static
+        # width and there really are surplus columns to get wrong.
+        loss_mask = self._loss_mask((6, 6))
+        torch.manual_seed(7)
+        legacy_anchors, legacy_keep = _legacy_sample_anchor_positions(
+            model, SEQ_LEN, loss_mask, self.DEVICE
+        )
+        torch.manual_seed(7)
+        anchors, keep = model._sample_anchor_positions(SEQ_LEN, loss_mask, self.DEVICE)
+        assert legacy_keep.shape[1] < keep.shape[1], "fixture no longer exercises truncation"
+        kept = anchors[keep]
+        assert torch.equal(kept, legacy_anchors[legacy_keep])
+
+    def test_shape_is_independent_of_batch_contents(self):
+        """The whole point: one shape, therefore one compile."""
+        model = self._model(num_anchors=8)
+        shapes = {
+            model._sample_anchor_positions(SEQ_LEN, self._loss_mask(lengths), self.DEVICE)[0].shape
+            for lengths in [(13, 13), (13, 5), (9, 9), (5, 3), (2, 1), (1, 1), (0, 0)]
+        }
+        assert len(shapes) == 1, f"n_blocks still varies with the batch: {shapes}"
+
+    def test_shape_is_min_of_num_anchors_and_sequence(self):
+        anchors, _ = self._model(num_anchors=4)._sample_anchor_positions(
+            SEQ_LEN, self._loss_mask((13, 13)), self.DEVICE
+        )
+        assert anchors.shape[1] == 4, "num_anchors should bind here"
+        anchors, _ = self._model(num_anchors=512)._sample_anchor_positions(
+            SEQ_LEN, self._loss_mask((13, 13)), self.DEVICE
+        )
+        assert anchors.shape[1] == SEQ_LEN - BLOCK_SIZE + 1, "the sequence should bind here"
+
+    def test_sampling_does_not_sync_on_the_batch(self):
+        """The old bound read valid_counts back to the host, stalling the pipeline.
+
+        Parsed rather than grepped: the prose explaining the removal names the call it
+        removed, so a substring search over the source matches its own docstring.
+        """
+        import ast
+        import inspect
+        import textwrap
+
+        tree = ast.parse(
+            textwrap.dedent(inspect.getsource(hf_dflash.HFDFlashModel._sample_anchor_positions))
+        )
+        called = {
+            node.func.attr
+            for node in ast.walk(tree)
+            if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)
+        }
+        assert not called & {"item", "tolist"}, f"host sync in anchor sampling: {called}"
+
+    def test_trailing_padding_blocks_do_not_change_the_loss(self):
+        """The padding the static shape introduces is weightless in the loss and accuracy.
+
+        Mathematically exact -- block_keep_mask zeroes those rows in both the numerator
+        and the normalizer -- but the reduction is over more elements, so compare at
+        float tolerance rather than bitwise.
+        """
+        model = self._model(num_anchors=8)
+        vocab, bsz, n_blocks, pad = 32, 1, 2, 3
+        torch.manual_seed(0)
+        input_ids = torch.randint(0, vocab, (bsz, SEQ_LEN))
+        loss_mask = torch.ones(bsz, SEQ_LEN)
+        logits = torch.randn(bsz, n_blocks * BLOCK_SIZE, vocab)
+        anchors = torch.tensor([[0, BLOCK_SIZE]])[:, :n_blocks]
+        keep = torch.ones(bsz, n_blocks)
+
+        base_loss, base_acc = model._compute_loss(logits, input_ids, anchors, keep, loss_mask)
+        padded_loss, padded_acc = model._compute_loss(
+            torch.cat([logits, torch.randn(bsz, pad * BLOCK_SIZE, vocab)], dim=1),
+            input_ids,
+            torch.cat([anchors, torch.zeros(bsz, pad, dtype=anchors.dtype)], dim=1),
+            torch.cat([keep, torch.zeros(bsz, pad)], dim=1),
+            loss_mask,
+        )
+        torch.testing.assert_close(padded_loss, base_loss)
+        assert padded_acc == pytest.approx(base_acc)
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py
index 6a0f884d33f..d33071ad128 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py
@@ -34,7 +34,7 @@
 import modelopt.torch.speculative as mtsp
 from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG
 from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel
-from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel
+from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel, _tvd_per_token
 from modelopt.torch.speculative.plugins.modeling_dflash import DFlashModule
 from modelopt.torch.speculative.plugins.modeling_dspark import DSparkModule
 
@@ -311,3 +311,216 @@ def test_export_config_has_dspark_fields(self, tmp_path):
         assert dc["shift_label"] is True
         assert "mask_token_id" in dc
         assert "target_layer_ids" in dc
+
+
+class TestTvdPerTokenChunking:
+    """``_tvd_per_token`` must be chunk-size-invariant, and must chunk via ``split``.
+
+    No other DSpark test reaches a second chunk: the tiny fixture yields N = 2 * 12 * 4
+    = 96 rows against the default ``chunk_size=1024``, so every existing test runs in a
+    single chunk and any chunk-boundary bug -- a mis-ordered ``cat``, a ragged tail
+    handled wrong, a row/label misalignment -- passes CI silently. These tests call the
+    helper directly so they can force many chunks on a tiny tensor.
+    """
+
+    @staticmethod
+    def _run(chunk_size, n=12, vocab=32):
+        """Forward + backward through ``_tvd_per_token`` from a fixed seed."""
+        torch.manual_seed(0)
+        final = torch.randn(n, vocab, requires_grad=True)
+        teacher = torch.randn(n, vocab)
+        out = _tvd_per_token(final, teacher, chunk_size=chunk_size)
+        # Weight the rows unequally, so a cat that reassembles the chunks in the wrong
+        # order cannot cancel out in the reduction.
+        (out * torch.arange(1, n + 1, dtype=out.dtype)).sum().backward()
+        return out.detach(), final.grad.detach()
+
+    # 12 rows: 5 and 7 leave a ragged last chunk (5+5+2, 7+5); 13 and 1024 exceed n.
+    @pytest.mark.parametrize("chunk_size", [1, 2, 3, 5, 7, 11, 12, 13, 1024])
+    def test_chunk_size_invariant(self, chunk_size):
+        ref_out, ref_grad = self._run(12)  # one chunk == the unchunked reference
+        out, grad = self._run(chunk_size)
+        assert torch.equal(out, ref_out), f"TVD value changed at chunk_size={chunk_size}"
+        assert torch.equal(grad, ref_grad), f"TVD grad changed at chunk_size={chunk_size}"
+
+    def test_chunks_via_split_not_slice(self):
+        """Pin the optimization itself, not just its result.
+
+        Slicing and splitting agree on every value, so no numerical test can tell them
+        apart -- only the backward graph can. Slicing costs O(n_chunks * N * vocab)
+        because each ``SliceBackward0`` zero-fills a full [N, vocab] tensor; at the
+        Gemma-4 shape that was 93.2 ms/step. This test is what stops the loop being
+        "simplified" back to ``final_logits[i : i + chunk_size]``.
+        """
+        final = torch.randn(8, 4, requires_grad=True)
+        teacher = torch.randn(8, 4)
+        out = _tvd_per_token(final, teacher, chunk_size=2)
+
+        # `alive` is load-bearing, not debris: accessing `.next_functions` hands back a
+        # FRESH python wrapper for each node every time, so a node we do not hold a
+        # reference to is freed the moment we pop it -- and CPython happily reuses that
+        # address for a later, different node. Without `alive` the id() check reports a
+        # false "already visited" and the walk truncates after three nodes, missing the
+        # Split/Slice node entirely (verified: both variants returned the same
+        # ['AbsBackward0', 'CatBackward0', 'SumBackward1']).
+        seen, visited, alive, stack = set(), set(), [], [out.grad_fn]
+        while stack:
+            fn = stack.pop()
+            if fn is None or id(fn) in visited:
+                continue
+            visited.add(id(fn))
+            alive.append(fn)
+            seen.add(type(fn).__name__)
+            stack.extend(nxt for nxt, _ in fn.next_functions)
+
+        assert any(name.startswith("SplitBackward") for name in seen), (
+            f"expected a SplitBackward node in the graph, saw: {sorted(seen)}"
+        )
+        assert not any(name.startswith("SliceBackward") for name in seen), (
+            f"chunking regressed to per-chunk slicing, saw: {sorted(seen)}"
+        )
+
+    def test_no_grad_path_matches_grad_path(self):
+        """The ``requires_grad=False`` branch skips checkpointing; values must not move."""
+        torch.manual_seed(0)
+        final = torch.randn(12, 32)
+        teacher = torch.randn(12, 32)
+        with torch.no_grad():
+            plain = _tvd_per_token(final, teacher, chunk_size=5)
+        grad_out = _tvd_per_token(final.requires_grad_(True), teacher, chunk_size=5)
+        assert torch.equal(plain, grad_out.detach())
+
+
+def _draft_args(model, n_blocks=2, bsz=1):
+    """Build the draft module's inputs the way the training forward does.
+
+    Worth going through the model's own helpers: the draft attends over the context as
+    keys but only its own blocks as queries, so target_hidden is seq_len long while
+    noise_embedding is n_blocks * block_size, and position_ids spans both. Hand-rolled
+    shapes silently disagree inside apply_rotary_pos_emb.
+    """
+    m = model.dflash_module
+    dt = m.fc.weight.dtype  # the draft carries the base model's dtype, not fp32
+    torch.manual_seed(0)
+    input_ids = torch.randint(1, model.dflash_config.vocab_size, (bsz, SEQ_LEN))
+    anchors = (torch.arange(n_blocks).unsqueeze(0) * BLOCK_SIZE).expand(bsz, -1).contiguous()
+    keep = torch.ones(bsz, n_blocks, dtype=torch.bool)
+    noise = model._build_noise_embedding(input_ids, anchors, keep, n_blocks).to(dt)
+    target = torch.randn(bsz, SEQ_LEN, m.fc.in_features, dtype=dt)
+    pos = model._build_position_ids(SEQ_LEN, anchors, input_ids.device)
+    mask = model._build_draft_attention_mask(
+        SEQ_LEN, anchors, keep, n_blocks, dt, input_ids.device, window=None
+    )
+    return input_ids, anchors, keep, (noise, target, pos, mask)
+
+
+def _dspark_model(use_compile=False, **cfg_kwargs):
+    model = get_tiny_llama(num_hidden_layers=4)
+    cfg = _get_dspark_config(**cfg_kwargs)
+    cfg["dflash_use_torch_compile"] = use_compile
+    mtsp.convert(model, [("dflash", cfg)])
+    model.train()
+    return model
+
+
+class TestDraftStackCompile:
+    """The draft stack is Inductor-compiled only when asked for, and only while training."""
+
+    def test_flag_off_keeps_the_eager_body(self):
+        m = _dspark_model(use_compile=False).dflash_module
+        assert m._body() == m._forward_body
+
+    def test_eval_keeps_the_eager_body(self):
+        """Generation runs this module at a varying length; under dynamic=False that would
+        mint one compile per length -- the cost the pinned block count exists to avoid."""
+        m = _dspark_model(use_compile=True).dflash_module
+        m.eval()
+        assert m._body() == m._forward_body
+
+    def test_compiled_body_is_built_once(self):
+        m = _dspark_model(use_compile=True).dflash_module
+        first = m._body()
+        assert first != m._forward_body
+        assert m._body() is first, "recompiled on a later step"
+
+    def test_compiled_matches_eager(self):
+        """Same weights, same inputs, both bodies -- same hidden states."""
+        model = _dspark_model(use_compile=True)
+        m = model.dflash_module
+        _, _, _, args = _draft_args(model)
+        m._maybe_init_rotary_emb(device=args[0].device)  # normally done by forward()
+
+        with torch.no_grad():
+            ref = m._forward_body(*args)
+            got = m._body()(*args)
+        assert m._body() != m._forward_body, "fixture did not actually compile"
+        # Inductor may fuse and reassociate, so this is agreement to the dtype's precision,
+        # not bitwise equality.
+        torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2)
+
+
+class TestDdpGradientCoverage:
+    """Every draft parameter must get a gradient on every batch, degenerate ones included.
+
+    This is the precondition for ddp_find_unused_parameters=false: DDP aborts the run the
+    first time a parameter that joined the reduction produces no gradient. The confidence
+    head is the one at risk -- it hangs off its own projection and never appears in
+    final_logits, so the branches that skip the loss terms have to reach it deliberately.
+    """
+
+    @staticmethod
+    def _model():
+        return _dspark_model(use_confidence_head=True, confidence_alpha=1.0)
+
+    @staticmethod
+    def _ungraded(model):
+        return [
+            n
+            for n, p in model.dflash_module.named_parameters()
+            if p.requires_grad and p.grad is None
+        ]
+
+    def test_batch_with_no_valid_anchor(self):
+        """Nothing is a training target, so forward() returns before building the draft."""
+        model = self._model()
+        torch.manual_seed(0)
+        input_ids = torch.randint(1, model.dflash_config.vocab_size, (2, SEQ_LEN))
+        out = model(
+            input_ids=input_ids,
+            attention_mask=torch.ones_like(input_ids),
+            labels=torch.full_like(input_ids, -100),
+        )
+        out.loss.backward()
+        assert not self._ungraded(model), f"no gradient for {self._ungraded(model)}"
+
+    def test_loss_branch_with_zero_total_weight(self):
+        """Anchors exist but every label position is masked, so the three terms are skipped.
+
+        Driven through the real module rather than synthetic logits: a gradient assertion
+        is only meaningful if the graph actually reaches the parameters.
+        """
+        model = self._model()
+        m = model.dflash_module
+        n_blocks, bsz = 2, 1
+        input_ids, anchors, _, args = _draft_args(model, n_blocks=n_blocks, bsz=bsz)
+        hidden = m(*args)
+        vocab = model.dflash_config.vocab_size
+        backbone_logits = torch.randn(bsz, n_blocks * BLOCK_SIZE, vocab, dtype=hidden.dtype)
+        # The teacher distribution is gathered before the degenerate branch is reached, so
+        # it has to be present even though this batch contributes nothing to the loss.
+        target_model_logits = torch.randn(bsz, SEQ_LEN, vocab, dtype=hidden.dtype)
+        final_logits, confidence_logits = model._apply_markov_head(
+            hidden, backbone_logits, input_ids, anchors, n_blocks
+        )
+        loss, _, _ = model._compute_dspark_loss(
+            backbone_logits,
+            final_logits,
+            confidence_logits,
+            input_ids,
+            anchors,
+            torch.ones(bsz, n_blocks),  # blocks are kept ...
+            torch.zeros(bsz, SEQ_LEN),  # ... but no label position carries weight
+            target_model_logits=target_model_logits,
+        )
+        loss.backward()
+        assert not self._ungraded(model), f"no gradient for {self._ungraded(model)}"
diff --git a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py
index efd77267b3c..43d641bbf39 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py
@@ -430,6 +430,88 @@ def test_lapped_slot_is_treated_as_miss(monkeypatch):
         ds[0]
 
 
+def _done_raising_handler(seq, n_layers, hidden, *, fail_first_n, calls):
+    """Sidecar handler whose /done raises a transport error for its first ``fail_first_n``
+    calls, then behaves normally. ``calls`` accumulates the paths hit."""
+    inner = _rdma_sidecar_handler(seq, n_layers, hidden)
+    state = {"done": 0}
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        calls.append(request.url.path)
+        if request.url.path == "/done":
+            state["done"] += 1
+            if state["done"] <= fail_first_n:
+                raise httpx.ConnectError("simulated sidecar failure")
+        return inner(request)
+
+    return handler
+
+
+def test_failed_done_is_treated_as_a_miss(monkeypatch):
+    """A /done that errors must not be read as "the slot is still ours".
+
+    The gen check /done performs is the only thing between a lapped ring slot and training
+    on another prompt's activations: the token_ids comparison that follows compares the
+    server's per-request record, not the bytes just read, so it passes on a mis-slotted
+    read. Trusting the read when /done is unreachable trades a visible resample for silent
+    corruption, so the failure must fail closed like an explicit lap does.
+    """
+    seq, n_layers, hidden = 8, 3, 16
+    calls: list[str] = []
+    _mock_rdma(
+        monkeypatch,
+        _done_raising_handler(seq, n_layers, hidden, fail_first_n=10**6, calls=calls),
+    )
+
+    ds = EagleVllmStreamingDataset(
+        entries=[{"conversation_id": "c-0", "messages": [{"role": "user", "content": "x"}]}],
+        tokenizer=_tokenizer_returning(seq),
+        config=EagleVllmStreamingConfig(
+            server_urls="http://mock:8000",
+            model="mock-model",
+            max_seq_len=seq,
+            fail_after_consecutive_skips=100,
+        ),
+    )
+    with pytest.raises(RuntimeError, match="no fetchable sample"):
+        ds[0]
+    assert "/done" in calls, "fixture never reached /done"
+
+
+def test_failed_done_resamples_instead_of_returning_the_read(monkeypatch):
+    """The discard is a resample, not a hard failure: the next entry is fetched and returned.
+
+    Pins that failing closed costs one extra round trip rather than breaking the epoch --
+    the cheap half of the asymmetry that justifies discarding.
+    """
+    seq, n_layers, hidden = 8, 3, 16
+    calls: list[str] = []
+    _mock_rdma(
+        monkeypatch,
+        _done_raising_handler(seq, n_layers, hidden, fail_first_n=1, calls=calls),
+    )
+
+    ds = EagleVllmStreamingDataset(
+        entries=[
+            {"conversation_id": f"c-{i}", "messages": [{"role": "user", "content": "x"}]}
+            for i in range(2)
+        ],
+        tokenizer=_tokenizer_returning(seq),
+        config=EagleVllmStreamingConfig(
+            server_urls="http://mock:8000",
+            model="mock-model",
+            max_seq_len=seq,
+            fail_after_consecutive_skips=100,
+        ),
+    )
+
+    batch = ds[0]
+    assert batch["base_model_hidden_states"].shape == (seq, hidden)
+    # Two prompts posted: the first read was thrown away, the second is what came back.
+    assert calls.count("/v1/completions") == 2
+    assert calls.count("/done") == 2
+
+
 def test_oversize_server_response_raises(monkeypatch):
     """If the server captured more tokens than max_seq_len (its connector max_tokens >
     our recv buffer), reading would silently truncate the slice; fail loud instead so the