diff --git a/examples/speculative_decoding/assets/dspark_nemotron35_warmstart_curves.png b/examples/speculative_decoding/assets/dspark_nemotron35_warmstart_curves.png
new file mode 100644
index 00000000000..b598fba1c87
Binary files /dev/null and b/examples/speculative_decoding/assets/dspark_nemotron35_warmstart_curves.png differ
diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py
index 2ba2635fbca..0691ca06f52 100644
--- a/examples/speculative_decoding/eagle_utils.py
+++ b/examples/speculative_decoding/eagle_utils.py
@@ -59,12 +59,16 @@ def make_speculative_data_module(
     train_len=None,
     answer_only_loss=False,
     shift_labels=True,
+    final_aux_is_base_hidden=False,
 ) -> dict:
     """Create data module for speculative decoding training.
 
     Args:
         shift_labels: If True, labels are shifted by 1 for autoregressive training (EAGLE3).
             If False, labels are unshifted for diffusion-style training (DFlash).
+        final_aux_is_base_hidden: Streaming only. True when the draft's top aux layer is the
+            base's final layer, so the last captured plane is both the final aux feature and
+            the base (KD-target) hidden instead of an extra dedicated plane.
     """
     # Load chat template from file if provided
     chat_template = None
@@ -115,6 +119,7 @@ def make_speculative_data_module(
             model=data_args.streaming_model_name,
             max_seq_len=train_len,
             answer_only_loss=answer_only_loss,
+            final_aux_is_base_hidden=final_aux_is_base_hidden,
         )
         train_dataset = EagleVllmStreamingDataset(
             entries=ds,
diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py
index ce9460637ac..8476d3d13c5 100644
--- a/examples/speculative_decoding/main.py
+++ b/examples/speculative_decoding/main.py
@@ -285,12 +285,41 @@ def train():
 
     print_rank_0("Loading dataset...")
     is_dflash = isinstance(recipe, ModelOptDFlashRecipe)
+    # A draft whose top aux layer already is the base's last layer (e.g. the released
+    # Nemotron-3.5 DSpark draft: aux ids [2,6,20,30,42,52] on a 52-layer base) cannot get a
+    # distinct extra capture plane for the base hidden — vLLM captures each layer once — so
+    # the streaming dataset must reuse the final plane for both roles. Derived from the
+    # model rather than configured by hand: it is a property of the draft, and a wrong
+    # manual value fails as a confusing matmul shape error deep in the draft's `fc`.
+    # Read the base depth the same way HFDFlashModel.modify does: offline/streaming loads
+    # the base with num_hidden_layers=0 (or a fake base) and stashes the real count in
+    # num_orig_hidden_layers, so num_hidden_layers alone would compare against 0 here.
+    _base_cfg = (
+        getattr(model.config, "text_config", None)
+        or getattr(model.config, "llm_config", None)
+        or model.config
+    )
+    _base_depth = getattr(_base_cfg, "num_orig_hidden_layers", None) or getattr(
+        _base_cfg, "num_hidden_layers", 0
+    )
+    final_aux_is_base_hidden = bool(
+        is_dflash
+        and getattr(model, "target_layer_ids", None)
+        and _base_depth
+        and max(model.target_layer_ids) + 1 >= _base_depth
+    )
+    print_rank_0(
+        f"Streaming plane layout: base_depth={_base_depth}, "
+        f"draft target_layer_ids={getattr(model, 'target_layer_ids', None)}, "
+        f"final_aux_is_base_hidden={final_aux_is_base_hidden}"
+    )
     data_module = make_speculative_data_module(
         tokenizer,
         recipe.data,
         train_len=training_args.training_seq_len,
         answer_only_loss=training_args.answer_only_loss,
         shift_labels=not is_dflash,
+        final_aux_is_base_hidden=final_aux_is_base_hidden,
     )
 
     callbacks = [EagleTrainingPlot(training_args.ar_validate_steps, training_args.estimate_ar)]
diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py
index 255b1d9ab04..fad15abc2db 100644
--- a/modelopt/torch/export/plugins/hf_spec_export.py
+++ b/modelopt/torch/export/plugins/hf_spec_export.py
@@ -412,10 +412,10 @@ def _export_config(self):
         else:
             config["layer_types"] = ["full_attention"] * draft_config.num_hidden_layers
 
-        # Sliding-window attention: all draft layers use non-causal SWA (MiMo-style). vLLM's
+        # Sliding-window attention: all draft layers use SWA. vLLM's
         # _resolve_layer_attention reads dflash_config.use_swa + swa_window_size; with
-        # layer_types left all "full_attention" it applies a non-causal sliding window to
-        # every draft layer (window from swa_window_size / top-level sliding_window).
+        # layer_types left all "full_attention" it applies a sliding window to every draft
+        # layer (window from swa_window_size / top-level sliding_window).
         swa_window = getattr(self.model, "dflash_swa_window_size", None)
         if swa_window is not None:
             config["sliding_window"] = swa_window
@@ -423,10 +423,23 @@ def _export_config(self):
                 {
                     "use_swa": True,
                     "swa_window_size": swa_window,
-                    "causal": False,
                 }
             )
 
+        # Block-internal attention pattern. Emitted unconditionally (not just under SWA):
+        # vLLM's _dflash_layer_causal treats dflash_config.causal as an all-layer override,
+        # and its default differs per layer type, so writing it explicitly is what keeps
+        # inference consistent with how the draft was actually trained.
+        config["dflash_config"]["causal"] = (
+            getattr(self.model, "dflash_draft_attention", "bidirectional") == "causal"
+        )
+
+        # Learnable per-head attention sink. vLLM reads dflash_config.attention_sink_bias to
+        # decide whether to build the sink parameter and pass it to its attention kernel.
+        if getattr(self.model, "dflash_attention_sink", False):
+            config["dflash_config"]["attention_sink_bias"] = True
+            config["attention_sink_bias"] = True
+
         # Inject the export-time YaRN rope_scaling from the dflash_export_rope_scaling
         # config field (empty dict disables). Mirrors eagle's eagle_export_rope_scaling.
         export_rope_scaling = getattr(self.model, "dflash_export_rope_scaling", None)
diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py
index 4df57d3035f..08e93b6f6a7 100644
--- a/modelopt/torch/speculative/config.py
+++ b/modelopt/torch/speculative/config.py
@@ -153,12 +153,64 @@ class DFlashConfig(ModeloptBaseConfig):
         default=None,
         description=(
             "Sliding-window attention (SWA) window size for the DFlash draft. When set, ALL "
-            "draft layers use non-causal sliding-window attention (MiMo-style): each draft "
-            "query attends only to context positions within `dflash_swa_window_size` tokens "
-            "before it, while block-internal attention stays bidirectional. None (default) "
-            "keeps full attention over all context. Must be >= dflash_block_size. Exported to "
-            "the draft config as dflash_config.use_swa/swa_window_size (+ top-level "
-            "sliding_window) so vLLM applies the same window at inference."
+            "draft layers use sliding-window attention: each draft query attends only to "
+            "context positions within `dflash_swa_window_size` tokens before it. None "
+            "(default) keeps full attention over all context. Must be >= dflash_block_size. "
+            "Exported to the draft config as dflash_config.use_swa/swa_window_size (+ "
+            "top-level sliding_window) so vLLM applies the same window at inference. Whether "
+            "block-internal attention is bidirectional or causal is controlled separately by "
+            "`dflash_draft_attention`."
+        ),
+    )
+
+    dflash_draft_attention: Literal["bidirectional", "causal"] = ModeloptField(
+        default="bidirectional",
+        description=(
+            "Attention pattern *inside* each draft block (context attention is always "
+            "restricted to positions before the block's anchor, and additionally windowed "
+            "when dflash_swa_window_size is set).\n"
+            "- 'bidirectional' (default): every query in a block sees all block_size draft "
+            "positions, including ones after it (MiMo-style). This is what ModelOpt has "
+            "always trained and matches drafts such as XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash "
+            "and z-lab/Qwen3.5-9B-DFlash.\n"
+            "- 'causal': a query at block position i only sees draft positions <= i, so the "
+            "block is predicted autoregressively. Required to faithfully train drafts whose "
+            "config declares dflash_config.causal=true, e.g. "
+            "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark.\n"
+            "Exported verbatim to dflash_config.causal, which vLLM's "
+            "qwen3_dflash._dflash_layer_causal reads as a per-model override."
+        ),
+    )
+
+    dflash_attention_sink: bool = ModeloptField(
+        default=False,
+        description=(
+            "Add a learnable per-head attention sink to every draft attention layer. The "
+            "sink is one extra logit per head appended to the attention logits before the "
+            "softmax and dropped afterwards, letting a head place probability mass nowhere "
+            "instead of being forced to attend within a (possibly short) window — the "
+            "GPT-OSS/Nemotron formulation. Adds one `self_attn.attention_sink_bias` "
+            "parameter of shape [num_attention_heads] per layer. Required to load and "
+            "continue training drafts whose checkpoint carries those weights, e.g. "
+            "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark. Exported to "
+            "dflash_config.attention_sink_bias for vLLM."
+        ),
+    )
+
+    dflash_init_checkpoint: str | None = ModeloptField(
+        default=None,
+        description=(
+            "Path to an exported draft checkpoint to warm-start from, so training continues "
+            "from published weights instead of a fresh random init. Accepts either a "
+            "directory in the deployment layout this repo exports (``model.safetensors`` "
+            "with no ``dflash_module.`` prefix, alongside ``config.json``) or the "
+            "``model.safetensors`` file itself. Weights are loaded into the draft module "
+            "after it is built, so the architecture still comes from "
+            "``dflash_architecture_config`` — the checkpoint must match it. Any mismatch "
+            "(missing, unexpected, or wrong-shaped tensors) raises rather than silently "
+            "leaving part of the draft randomly initialized. ``embed_tokens``/``lm_head`` "
+            "entries are ignored: the draft takes those from the base model. None "
+            "(default) trains from scratch."
         ),
     )
 
@@ -235,8 +287,8 @@ def _check_dpace_alpha(self) -> "DFlashConfig":
         if not 0.0 < self.dflash_dpace_alpha <= 1.0:
             raise ValueError(f"dflash_dpace_alpha must be in (0, 1], got {self.dflash_dpace_alpha}")
         if self.dflash_swa_window_size is not None:
-            # Block-internal attention is left un-windowed (bidirectional), so the window must
-            # cover a full block; otherwise the effective inference window would differ.
+            # Block-internal attention is left un-windowed, so the window must cover a full
+            # block; otherwise the effective inference window would differ.
             if self.dflash_swa_window_size < self.dflash_block_size:
                 raise ValueError(
                     f"dflash_swa_window_size ({self.dflash_swa_window_size}) must be >= "
diff --git a/modelopt/torch/speculative/dflash/dflash_model.py b/modelopt/torch/speculative/dflash/dflash_model.py
index 24f2143bf84..3ce06afeda8 100644
--- a/modelopt/torch/speculative/dflash/dflash_model.py
+++ b/modelopt/torch/speculative/dflash/dflash_model.py
@@ -49,4 +49,7 @@ 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_draft_attention = config.dflash_draft_attention
+        self.dflash_attention_sink = config.dflash_attention_sink
+        self.dflash_init_checkpoint = config.dflash_init_checkpoint
         self.dflash_export_rope_scaling = config.dflash_export_rope_scaling
diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index e0d63bde136..616491dcf32 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -72,6 +72,7 @@
 """
 
 import logging
+from pathlib import Path
 from typing import Any
 
 import torch
@@ -441,15 +442,40 @@ def modify(self, config):
             self.dflash_config.hidden_size // self.dflash_config.num_attention_heads,
         )
         self.dflash_config.block_size = self.dflash_block_size
-
-        # Target layer IDs
+        # Attention-shape knobs the draft module needs at build time. Carried on the draft
+        # config (not read from DFlashConfig) so _build_draft_module stays a pure function
+        # of dflash_config and restored checkpoints rebuild an identical module.
+        self.dflash_config.attention_sink_bias = self.dflash_attention_sink
+
+        # Target layer IDs: which base layers feed the draft's feature-fusion `fc`.
+        # An explicit dflash_architecture_config.target_layer_ids wins — a published draft
+        # is trained against specific capture points (e.g. the Nemotron-3.5 DSpark draft's
+        # [1,5,19,29,41,51]), and recomputing the uniform default would both mis-shape `fc`
+        # and feed the draft features from layers it never saw. Otherwise derive the
+        # uniform default.
         num_target_layers = (
             base_config.num_orig_hidden_layers
             if self.dflash_offline
             else base_config.num_hidden_layers
         )
         num_draft_layers = self.dflash_config.num_hidden_layers
-        self.target_layer_ids = build_target_layer_ids(num_target_layers, num_draft_layers)
+        user_target_layer_ids = config.dflash_architecture_config.get("target_layer_ids")
+        if user_target_layer_ids:
+            if len(user_target_layer_ids) != num_draft_layers:
+                raise ValueError(
+                    f"dflash_architecture_config.target_layer_ids has "
+                    f"{len(user_target_layer_ids)} entries but the draft has "
+                    f"{num_draft_layers} layers; one target layer per draft layer is required."
+                )
+            if max(user_target_layer_ids) >= num_target_layers:
+                raise ValueError(
+                    f"dflash_architecture_config.target_layer_ids {user_target_layer_ids} "
+                    f"references a layer beyond the base model's {num_target_layers} layers."
+                )
+            self.target_layer_ids = list(user_target_layer_ids)
+            logger.info("DFlash: using explicit target_layer_ids %s", self.target_layer_ids)
+        else:
+            self.target_layer_ids = build_target_layer_ids(num_target_layers, num_draft_layers)
         self.dflash_config.target_layer_ids = self.target_layer_ids
 
         # mask_token_id: validated by DFlashConfig, auto-detected from tokenizer context
@@ -466,6 +492,10 @@ def modify(self, config):
         # Factory hook: subclasses (e.g. Domino) override to build an augmented
         # draft module while reusing all of DFlash's modify() setup.
         self.dflash_module = self._build_draft_module(self.dflash_config)
+        # Warm start from an exported draft checkpoint, before the dtype/device move below
+        # so the loaded tensors get cast alongside the rest of the module.
+        if self.dflash_init_checkpoint:
+            self._load_init_checkpoint(self.dflash_init_checkpoint)
         # Match base model dtype/device. Skip if base is on meta (during from_pretrained
         # restore — the model will be moved to the correct device after weight loading).
         if self.dflash_offline:
@@ -486,6 +516,81 @@ def _build_draft_module(self, dflash_config):
         """Build the draft module. Subclasses override to use an augmented module."""
         return DFlashModule(dflash_config)
 
+    # Draft-module entries that legitimately come from the base model rather than the
+    # exported draft checkpoint, so their absence (or presence) is not an error.
+    _INIT_CKPT_IGNORED_KEYS = ("embed_tokens.weight", "lm_head.weight")
+
+    def _load_init_checkpoint(self, path: str):
+        """Warm-start ``self.dflash_module`` from an exported draft checkpoint.
+
+        Accepts either the export directory (containing ``model.safetensors``) or the
+        safetensors file itself. The architecture is fixed by ``dflash_architecture_config``
+        at this point, so the checkpoint has to match it: any missing, unexpected, or
+        wrong-shaped tensor raises. Loading part of a draft and leaving the rest randomly
+        initialized looks like a warm start but trains from a corrupted starting point, so
+        it is rejected instead of warned about.
+        """
+        from safetensors.torch import load_file
+
+        ckpt = Path(path)
+        if ckpt.is_dir():
+            ckpt = ckpt / "model.safetensors"
+        if not ckpt.is_file():
+            raise FileNotFoundError(
+                f"dflash_init_checkpoint: no draft weights at {ckpt}. Expected an exported "
+                "draft directory containing model.safetensors, or the file itself."
+            )
+
+        state_dict = load_file(str(ckpt))
+        # Tolerate a `dflash_module.` prefix so a raw training checkpoint also works.
+        state_dict = {
+            (k.split("dflash_module.", 1)[1] if "dflash_module." in k else k): v
+            for k, v in state_dict.items()
+        }
+        state_dict = {
+            k: v
+            for k, v in state_dict.items()
+            if k not in self._INIT_CKPT_IGNORED_KEYS and "rotary_emb" not in k
+        }
+
+        # Shape-check against the module's own view of each key. Subclasses may remap keys
+        # on load (DSpark accepts a nested ``markov_head.`` layout), so resolve through the
+        # same hooks first — otherwise a wrong-shaped remapped tensor would skip this check
+        # and fail later with a much less obvious error.
+        module_sd = self.dflash_module.state_dict()
+        resolved = dict(state_dict)
+        for hook in self.dflash_module._load_state_dict_pre_hooks.values():
+            hook(resolved, "", None, True, [], [], [])
+        mismatched = [
+            f"{k}: checkpoint {tuple(v.shape)} vs module {tuple(module_sd[k].shape)}"
+            for k, v in resolved.items()
+            if k in module_sd and v.shape != module_sd[k].shape
+        ]
+        if mismatched:
+            raise ValueError(
+                "dflash_init_checkpoint: shape mismatch between "
+                f"{ckpt} and the configured draft architecture:\n  " + "\n  ".join(mismatched)
+            )
+
+        # strict=False, then check by hand: the module's own load hooks (e.g. DSpark's
+        # markov_head remap) run first, and buffers such as rotary_emb are excluded above.
+        incompatible = self.dflash_module.load_state_dict(state_dict, strict=False)
+        missing = [
+            k
+            for k in incompatible.missing_keys
+            if "rotary_emb" not in k and k not in self._INIT_CKPT_IGNORED_KEYS
+        ]
+        if missing or incompatible.unexpected_keys:
+            raise ValueError(
+                f"dflash_init_checkpoint: {ckpt} does not match the configured draft "
+                "architecture.\n"
+                f"  missing from checkpoint: {sorted(missing)}\n"
+                f"  unexpected in checkpoint: {sorted(incompatible.unexpected_keys)}"
+            )
+        logger.info(
+            "DFlash: warm-started draft module from %s (%d tensors).", ckpt, len(state_dict)
+        )
+
     def get_exporter(self):
         """Get the exporter for the DFlash draft model."""
         from modelopt.torch.export.plugins.hf_spec_export import DFlashExporter
@@ -567,13 +672,18 @@ def _build_position_ids(self, seq_len, anchor_positions, device):
     def _build_draft_attention_mask(
         self, seq_len, anchor_positions, block_keep_mask, n_blocks, dtype, device, window=None
     ):
-        """Build SDPA attention mask: context (causal) + draft (bidirectional within block).
-
-        When ``window`` is not None, all layers use non-causal sliding-window attention
-        (MiMo-style): each draft query only sees context positions within ``window`` tokens
-        before its own position. Block-internal attention stays bidirectional and is left
-        un-windowed (the config enforces ``window >= block_size``, so a full block always
-        fits inside the window and windowing it would be a no-op).
+        """Build SDPA attention mask: context (causal) + draft (per ``dflash_draft_attention``).
+
+        When ``window`` is not None, all layers use sliding-window attention: each draft
+        query only sees context positions within ``window`` tokens before its own position.
+        Block-internal attention is left un-windowed (the config enforces
+        ``window >= block_size``, so a full block always fits inside the window and windowing
+        it would be a no-op).
+
+        Block-internal visibility follows ``self.dflash_draft_attention``:
+        ``"bidirectional"`` (default, MiMo-style) lets every query see the whole block, while
+        ``"causal"`` restricts a query at block position ``i`` to draft positions ``<= i`` so
+        the block is modelled autoregressively.
         """
         bsz = anchor_positions.shape[0]
         block_size = self.dflash_block_size
@@ -598,6 +708,12 @@ def _build_draft_attention_mask(
         is_draft = kv_indices >= seq_len
         kv_block_ids = (kv_indices - seq_len) // block_size
         mask_draft = is_draft & (q_block_ids == kv_block_ids)
+        if self.dflash_draft_attention == "causal":
+            # Autoregressive within the block: query at block position i sees draft
+            # positions <= i only. Compare positions *within* the block so the term is
+            # independent of which block the query belongs to.
+            kv_pos_in_block = (kv_indices - seq_len) % block_size
+            mask_draft = mask_draft & (kv_pos_in_block <= (q_indices % block_size))
         # Valid block
         valid_block = block_keep_mask.view(bsz, 1, n_blocks, 1).repeat_interleave(block_size, dim=2)
 
@@ -609,24 +725,33 @@ def _build_draft_attention_mask(
         return attn_mask
 
     def _build_generate_swa_mask(self, ctx_len, bsz, dtype, device):
-        """Generation-time SWA mask [B, 1, block_size, ctx_len + block_size], or None.
+        """Generation-time mask [B, 1, block_size, ctx_len + block_size], or None.
 
-        Returns None with full attention (KV cache with no mask): all positions attend
-        freely to context and each other within the block. With sliding-window attention,
+        Returns None only when there is nothing to mask: full attention over the context
+        *and* bidirectional blocks (KV cache with no mask). With sliding-window attention,
         each block query only sees context within ``dflash_swa_window_size`` tokens before
         its real position (ctx_len + position-in-block), matching training and vLLM
-        inference; block kv stays fully visible (bidirectional / un-windowed).
+        inference; block kv is left un-windowed. With ``dflash_draft_attention="causal"``
+        the block is additionally lower-triangular, mirroring
+        :meth:`_build_draft_attention_mask` so generation matches training.
         """
-        if self.dflash_swa_window_size is None:
-            return None
         window = self.dflash_swa_window_size
+        causal = self.dflash_draft_attention == "causal"
+        if window is None and not causal:
+            return None
         block_size = self.dflash_block_size
         kv_len = ctx_len + block_size
         kv_idx = torch.arange(kv_len, device=device).view(1, 1, 1, -1)
-        q_real_pos = torch.arange(ctx_len, ctx_len + block_size, device=device).view(1, 1, -1, 1)
+        q_pos_in_block = torch.arange(block_size, device=device).view(1, 1, -1, 1)
+        q_real_pos = ctx_len + q_pos_in_block
         is_ctx = kv_idx < ctx_len
-        # Context kv kept iff within the window; block kv (>= ctx_len) always visible.
-        keep = (~is_ctx) | (kv_idx > q_real_pos - window)
+        # Context kv kept iff within the window (when windowing); block kv always visible
+        # unless the block is causal, in which case only positions <= the query's.
+        keep_ctx = is_ctx if window is None else (is_ctx & (kv_idx > q_real_pos - window))
+        keep_block = ~is_ctx
+        if causal:
+            keep_block = keep_block & ((kv_idx - ctx_len) <= q_pos_in_block)
+        keep = keep_ctx | keep_block
         attn_mask = torch.zeros(bsz, 1, block_size, kv_len, device=device, dtype=dtype)
         attn_mask.masked_fill_(~keep, torch.finfo(dtype).min)
         return attn_mask
diff --git a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py
index 0f9713948fd..ad6082eb71a 100644
--- a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py
+++ b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py
@@ -370,6 +370,15 @@ class EagleVllmStreamingConfig(StreamingConfig):
     # vLLM captures the residual stream BEFORE the final norm, so the trainer must re-apply it
     # before lm_head (see HFDFlashModel.forward). Set False for a post-norm producer.
     base_hidden_prenorm: bool = True
+    # Whether the LAST captured plane doubles as both the final aux feature and the base
+    # (KD/distillation-target) hidden. Normally they are distinct: the draft's aux layers
+    # sit below the base's last layer, so the producer captures ``aux + [final]`` and the
+    # trainer peels the extra plane off. A draft whose top aux id already IS the base's
+    # final layer (e.g.
+    # nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark, aux ids [2,6,20,30,42,52]
+    # on a 52-layer base) cannot supply a distinct extra id — vLLM captures each layer at
+    # most once — so the final plane is reused for both roles.
+    final_aux_is_base_hidden: bool = False
 
     @field_validator("server_urls", mode="before")
     @classmethod
@@ -579,8 +588,15 @@ def _format(self, fetched: EagleFetchPayload) -> dict[str, torch.Tensor]:
         hidden_states = fetched["hidden_states"]
         loss_mask = fetched["loss_mask"]
 
+        # The last plane is always the base (KD-target) hidden. It is normally an extra
+        # plane on top of the aux features; when the draft's top aux layer already is the
+        # base's final layer the producer cannot emit a distinct extra plane, so the same
+        # plane serves both roles (see ``final_aux_is_base_hidden``).
         base_model_hidden_states = hidden_states[:, -1, :]
-        aux_hidden_states = hidden_states[:, :-1, :].reshape(hidden_states.shape[0], -1)
+        aux_planes = (
+            hidden_states if self.config.final_aux_is_base_hidden else hidden_states[:, :-1, :]
+        )
+        aux_hidden_states = aux_planes.reshape(hidden_states.shape[0], -1)
 
         input_ids = token_ids.to(torch.int64)
         labels = torch.full_like(input_ids, IGNORE_TOKEN_ID)
diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py
index 6463cb4109d..3513ea53cc5 100644
--- a/modelopt/torch/speculative/plugins/modeling_dflash.py
+++ b/modelopt/torch/speculative/plugins/modeling_dflash.py
@@ -44,6 +44,7 @@
 from dataclasses import dataclass
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 from transformers.models.qwen3.modeling_qwen3 import Qwen3MLP as _MLP_CLS  # noqa: N814
@@ -51,6 +52,7 @@
 from transformers.models.qwen3.modeling_qwen3 import (
     Qwen3RotaryEmbedding as _ROTARY_CLS,  # noqa: N814
 )
+from transformers.models.qwen3.modeling_qwen3 import repeat_kv
 from transformers.models.qwen3.modeling_qwen3 import rotate_half as _rotate_half
 
 from .modeling_final_norm import _maybe_apply_base_final_norm
@@ -155,6 +157,16 @@ def __init__(self, config, layer_idx):
         self.q_norm = _NORM_CLS(self.head_dim, eps=config.rms_norm_eps)
         self.k_norm = _NORM_CLS(self.head_dim, eps=config.rms_norm_eps)
 
+        # Learnable per-head attention sink (GPT-OSS / Nemotron formulation): one extra
+        # logit per head, appended before the softmax and dropped after, so a head can put
+        # probability mass "nowhere" instead of being forced to attend inside its window.
+        # Named to match the deployed checkpoints' `self_attn.attention_sink_bias`.
+        self.attention_sink_bias = (
+            nn.Parameter(torch.zeros(self.num_heads))
+            if getattr(config, "attention_sink_bias", False)
+            else None
+        )
+
         # Resolve HF attention function
         self._attn_fn = None
         # Qwen3 uses sliding window attention on some layers (config.layer_types)
@@ -205,21 +217,70 @@ 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,
-        )
+        if self.attention_sink_bias is not None:
+            if self.sliding_window is not None:
+                # The eager sink path applies only the caller-supplied mask; a per-layer
+                # window from config.layer_types would be silently dropped. DFlash windows
+                # the context through the attention mask instead (dflash_swa_window_size),
+                # so this combination is rejected rather than trained with the wrong mask.
+                raise NotImplementedError(
+                    "dflash_attention_sink is not supported together with a per-layer "
+                    "sliding window from dflash_architecture_config.layer_types. Use "
+                    "dflash_swa_window_size for the draft's sliding window instead."
+                )
+            attn_output = self._sink_attention(q, k, v, attention_mask)
+        else:
+            # 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)
 
+    def _sink_attention(self, q, k, v, attention_mask):
+        """Eager attention with a learnable per-head sink logit.
+
+        The sink is an extra column appended to the attention logits before the softmax and
+        dropped immediately after, so it consumes probability mass without contributing to
+        the output. Fused SDPA/flash kernels cannot express that extra column, so this path
+        is eager; it runs only when ``dflash_attention_sink`` is enabled.
+
+        Mirrors ``transformers``' GPT-OSS ``eager_attention_forward``, including the
+        max-subtraction before the softmax that keeps bf16 training from overflowing.
+
+        Returns ``[B, q_len, num_heads, head_dim]`` to match the HF attention interface.
+        """
+        sink_bias = self.attention_sink_bias
+        assert sink_bias is not None, "_sink_attention requires dflash_attention_sink=True"
+
+        k = repeat_kv(k, self.num_key_value_groups)
+        v = repeat_kv(v, self.num_key_value_groups)
+
+        attn_weights = torch.matmul(q, k.transpose(2, 3)) * self.scaling
+        if attention_mask is not None:
+            # [B, 1, Q, KV] additive mask, already sliced to the kv length by the caller.
+            attn_weights = attn_weights + attention_mask[..., : k.shape[-2]]
+
+        sinks = sink_bias.view(1, -1, 1, 1).expand(
+            attn_weights.shape[0], -1, attn_weights.shape[-2], -1
+        )
+        combined = torch.cat([attn_weights, sinks.to(attn_weights.dtype)], dim=-1)
+        combined = combined - combined.amax(dim=-1, keepdim=True)
+        probs = F.softmax(combined, dim=-1, dtype=torch.float32).to(q.dtype)
+        attn_weights = probs[..., :-1]  # drop the sink column
+        attn_weights = F.dropout(
+            attn_weights, p=self.attention_dropout if self.training else 0.0, training=self.training
+        )
+        return torch.matmul(attn_weights, v).transpose(1, 2).contiguous()
+
 
 class DFlashDecoderLayer(nn.Module):
     """Draft decoder layer with KV injection."""
diff --git a/modelopt/torch/speculative/plugins/modeling_dspark.py b/modelopt/torch/speculative/plugins/modeling_dspark.py
index aa1a4181bbd..083af7b61ed 100644
--- a/modelopt/torch/speculative/plugins/modeling_dspark.py
+++ b/modelopt/torch/speculative/plugins/modeling_dspark.py
@@ -120,6 +120,30 @@ def __init__(self, config):
         # existed, so initialize the new layers explicitly.
         self._init_head_weights(config)
 
+        self._register_load_state_dict_pre_hook(self._remap_nested_head_keys, with_module=False)
+
+    # Head submodules that some published checkpoints nest under a `markov_head.` parent.
+    _NESTED_HEAD_PREFIX = "markov_head."
+
+    @classmethod
+    def _remap_nested_head_keys(cls, state_dict, prefix, *args, **kwargs):
+        """Accept head weights nested under ``markov_head.`` as well as flat.
+
+        ModelOpt keeps the head tensors flat on this module (``markov_w1`` /
+        ``markov_w2`` / ...), matching the upstream DeepSpec layout, and exports them that
+        way. Some released drafters — e.g.
+        ``nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark`` — instead nest them
+        under a ``markov_head.`` parent module. Rewriting those keys in place here lets
+        either layout load without changing the export format (which would break drafters
+        already trained and deployed with the flat names).
+        """
+        nested = prefix + cls._NESTED_HEAD_PREFIX
+        for key in [k for k in state_dict if k.startswith(nested)]:
+            flat = prefix + key[len(nested) :]
+            # A flat key already present wins: never clobber an explicit match.
+            state_dict.setdefault(flat, state_dict.pop(key))
+            state_dict.pop(key, None)
+
     def _init_head_weights(self, config):
         """Initialize the head Linear/Embedding layers (matching HF _init_weights std)."""
         std = getattr(config, "initializer_range", 0.02)
diff --git a/modelopt/torch/speculative/plugins/modeling_final_norm.py b/modelopt/torch/speculative/plugins/modeling_final_norm.py
index 718d591b662..a882ef4b084 100644
--- a/modelopt/torch/speculative/plugins/modeling_final_norm.py
+++ b/modelopt/torch/speculative/plugins/modeling_final_norm.py
@@ -90,6 +90,9 @@ def extra_repr(self):
     "deepseek_v3": "rmsnorm",
     "kimi_k2": "rmsnorm",  # Kimi-K2 / K2-Thinking (DeepSeek-V3 arch) report model_type "kimi_k2"
     "kimi_k25": "rmsnorm",  # Kimi-K2.5 / K2.6 / K2.7 all report model_type "kimi_k25"
+    # Nemotron-H hybrid Mamba/attention/MoE (e.g. NVIDIA-Nemotron-3.5-Lightning-30B-A3B).
+    # Its final norm (NemotronHModel.norm_f) is a plain RMSNorm despite the hybrid stack.
+    "nemotron_h": "rmsnorm",
     # M3's final norm is always gemma-style; map it here too so a config that lost its
     # use_gemma_norm flag still gets the correct flavor instead of silently dropping the +1.
     "minimax_m3_vl_text": "gemma_rmsnorm",
diff --git a/modelopt_recipes/general/speculative_decoding/dspark_nemotron35_warmstart.yaml b/modelopt_recipes/general/speculative_decoding/dspark_nemotron35_warmstart.yaml
new file mode 100644
index 00000000000..7366dee7b4a
--- /dev/null
+++ b/modelopt_recipes/general/speculative_decoding/dspark_nemotron35_warmstart.yaml
@@ -0,0 +1,102 @@
+# Warm-start fine-tune of the released Nemotron-3.5-Lightning DSpark drafter.
+#
+# Continues training `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark` from its
+# published weights, against `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16` served by
+# vLLM (data.mode=streaming).
+#
+# TODO: replace the hand-copied fields below with a proper checkpoint/config converter.
+# Everything under `dflash_architecture_config`, plus block_size / mask_token_id /
+# swa_window_size / draft_attention / attention_sink / target_layer_ids, is transcribed by
+# hand from the drafter's own config.json. Only the shape-bearing fields fail loudly when
+# mistyped; the rest (mask token, causal flag, window size, block size) train "successfully"
+# on the wrong setting and only show up later as a mysteriously low acceptance length. A
+# converter should derive this whole block from the checkpoint's config.json — including its
+# aliases (`pard_token` for mask_token_id, `dspark_markov_rank` for markov_rank,
+# `dflash_query_causal` for causal, top-level `sliding_window` / `attention_sink_bias`) —
+# and reconcile the weight layout (the release nests the Markov head under `markov_head.`,
+# which the loader remaps today).
+#
+# The base model's config.json also needs its layer-type vocabulary updated for vLLM's
+# transformers-5 path (`mamba`->`linear_attention`, `attention`->`full_attention`, plus a
+# matching `hybrid_override_pattern`); that conversion is not covered here either.
+
+metadata:
+  recipe_type: speculative_dflash
+  description: Warm-start fine-tune of the released Nemotron-3.5 DSpark drafter (streaming).
+
+model:
+  model_name_or_path: <path to NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16>
+  trust_remote_code: true
+  use_fake_base_for_offline: true
+
+data:
+  mode: streaming
+  data_path: <path to a jsonl corpus with real assistant turns>
+  offline_data_path:
+  # The stock Nemotron template has no {% generation %} tags, so answer_only_loss would get
+  # an all-zero assistant mask. Point this at a tagged copy.
+  chat_template: <path to nemotron chat template with {% generation %} tags>
+
+training:
+  output_dir: <output dir>
+  num_train_epochs: 1
+  per_device_train_batch_size: 1
+  gradient_accumulation_steps: 1
+  # Low by default: this continues from converged weights rather than training from scratch.
+  learning_rate: 1.0e-5
+  warmup_steps: 10
+  training_seq_len: 2048
+  logging_steps: 10
+  save_steps: 1000
+  cp_size: 1
+  dp_shard_size: 1
+  disable_tqdm: true
+  # Eval runs the DFlash backbone only (no Markov head), so AR would understate the model.
+  estimate_ar: false
+  ar_validate_steps: 0
+  answer_only_loss: true
+  do_eval: false
+  lr_scheduler_type: linear
+  save_strategy: steps
+  weight_decay: 0.0
+  max_grad_norm: 1.0
+  dataloader_drop_last: true
+  bf16: true
+  tf32: true
+  remove_unused_columns: false
+  ddp_find_unused_parameters: true
+  ddp_timeout: 1800
+  report_to: tensorboard
+
+dflash:
+  # --- transcribed from the released drafter's config.json (see TODO above) ---
+  dflash_init_checkpoint: <path to NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark>
+  dflash_block_size: 8
+  dflash_mask_token_id: 990
+  dflash_swa_window_size: 1024
+  dflash_draft_attention: causal
+  dflash_attention_sink: true
+  # --- training knobs ---
+  dflash_num_anchors: 64
+  dflash_use_torch_compile: false
+  dflash_self_logit_distillation: false
+  dflash_loss_objective: dpace
+  dflash_ce_loss_alpha: 0.1
+  dflash_l1_loss_alpha: 0.9
+  # The released checkpoint carries no confidence head.
+  dflash_confidence_head_alpha: 0.0
+  dflash_report_acc: true
+  dflash_architecture_config:
+    num_hidden_layers: 6
+    num_attention_heads: 32
+    num_key_value_heads: 2
+    head_dim: 128
+    intermediate_size: 6144
+    rms_norm_eps: 1.0e-6
+    projector_type: dspark
+    markov_rank: 512
+    markov_head_type: vanilla
+    use_confidence_head: false
+    # The layers whose hidden states feed `fc`. Must match the checkpoint: the uniform
+    # default for a 52-layer base would be [1,11,20,30,39,49], i.e. different layers.
+    target_layer_ids: [1, 5, 19, 29, 41, 51]
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py
index 6a0f884d33f..732898804ce 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py
@@ -29,13 +29,17 @@
 import pytest
 import torch
 from _test_utils.torch.transformers_models import get_tiny_llama
-from safetensors.torch import load_file
+from safetensors.torch import load_file, save_file
 
 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.modeling_dflash import DFlashModule
+from modelopt.torch.speculative.plugins.modeling_dflash import (
+    DFlashModule,
+    build_target_layer_ids,
+    repeat_kv,
+)
 from modelopt.torch.speculative.plugins.modeling_dspark import DSparkModule
 
 BLOCK_SIZE = 4
@@ -311,3 +315,410 @@ 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 TestDraftAttentionPattern:
+    """dflash_draft_attention selects the block-internal attention pattern."""
+
+    def _make_model(self, draft_attention, window=None, attention_sink=False):
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = _get_dspark_config()
+        config["dflash_draft_attention"] = draft_attention
+        config["dflash_attention_sink"] = attention_sink
+        if window is not None:
+            config["dflash_swa_window_size"] = window
+        mtsp.convert(model, [("dflash", config)])
+        return model
+
+    def _draft_block(self, model, seq_len=SEQ_LEN, n_blocks=2):
+        """Return the [block_size, block_size] draft-vs-draft visibility of block 0."""
+        anchors = torch.tensor([[5, 9]])[:, :n_blocks]
+        keep = torch.ones(1, n_blocks, dtype=torch.bool)
+        mask = model._build_draft_attention_mask(
+            seq_len,
+            anchors,
+            keep,
+            n_blocks,
+            torch.float32,
+            torch.device("cpu"),
+            window=model.dflash_swa_window_size,
+        )
+        # additive mask: 0 == visible, -inf == masked
+        visible = mask[0, 0] == 0
+        return visible[:BLOCK_SIZE, seq_len : seq_len + BLOCK_SIZE]
+
+    def test_default_is_bidirectional(self):
+        model = self._make_model("bidirectional")
+        assert model.dflash_draft_attention == "bidirectional"
+        assert self._draft_block(model).all(), "every query should see the whole block"
+
+    def test_causal_block_is_lower_triangular(self):
+        model = self._make_model("causal")
+        block = self._draft_block(model)
+        expected = torch.tril(torch.ones(BLOCK_SIZE, BLOCK_SIZE, dtype=torch.bool))
+        assert torch.equal(block, expected), f"expected lower-triangular, got {block}"
+
+    def test_causal_does_not_change_context_visibility(self):
+        """Only draft-vs-draft visibility changes; context masking is untouched."""
+        anchors = torch.tensor([[5, 9]])
+        keep = torch.ones(1, 2, dtype=torch.bool)
+        args = (SEQ_LEN, anchors, keep, 2, torch.float32, torch.device("cpu"))
+        bi = self._make_model("bidirectional")._build_draft_attention_mask(*args)
+        ca = self._make_model("causal")._build_draft_attention_mask(*args)
+        assert torch.equal(bi[..., :SEQ_LEN], ca[..., :SEQ_LEN])
+
+    def test_causal_generate_mask_is_lower_triangular(self):
+        """The generation-time mask matches training even without SWA."""
+        model = self._make_model("causal")
+        mask = model._build_generate_swa_mask(SEQ_LEN, 1, torch.float32, torch.device("cpu"))
+        assert mask is not None, "causal blocks need a mask even with full context attention"
+        visible = mask[0, 0] == 0
+        assert visible[:, :SEQ_LEN].all(), "full attention over context"
+        expected = torch.tril(torch.ones(BLOCK_SIZE, BLOCK_SIZE, dtype=torch.bool))
+        assert torch.equal(visible[:, SEQ_LEN:], expected)
+
+    def test_bidirectional_full_attention_needs_no_mask(self):
+        """Legacy behaviour: nothing to mask means no mask is built."""
+        model = self._make_model("bidirectional")
+        assert (
+            model._build_generate_swa_mask(SEQ_LEN, 1, torch.float32, torch.device("cpu")) is None
+        )
+
+    def test_causal_swa_trains_and_generates(self):
+        """End-to-end: causal + SWA produces a finite loss and drafts tokens."""
+        model = self._make_model("causal", window=6)
+        torch.manual_seed(0)
+        input_ids = torch.randint(1, model.dflash_config.vocab_size, (2, SEQ_LEN))
+        model.train()
+        out = model(
+            input_ids=input_ids,
+            attention_mask=torch.ones_like(input_ids),
+            labels=input_ids.clone(),
+        )
+        assert torch.isfinite(out.loss)
+        model.eval()
+        base_token, draft_tokens = model.pseudo_speculative_generate(input_ids[:1], steps=3)
+        assert base_token.shape == (1, 1)
+        assert draft_tokens.shape == (1, 3)
+
+    def test_export_records_draft_attention(self, tmp_path):
+        """dflash_config.causal is exported for both settings, with and without SWA."""
+        for mode, expected in (("bidirectional", False), ("causal", True)):
+            model = self._make_model(mode)
+            export_dir = tmp_path / f"exp_{mode}"
+            model.get_exporter().export(export_dir)
+            with open(export_dir / "config.json") as f:
+                cfg = json.load(f)
+            assert cfg["dflash_config"]["causal"] is expected
+
+
+class TestAttentionSink:
+    """dflash_attention_sink adds a learnable per-head sink to every draft layer."""
+
+    def _make_model(self, attention_sink=True, draft_attention="causal"):
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = _get_dspark_config()
+        config["dflash_attention_sink"] = attention_sink
+        config["dflash_draft_attention"] = draft_attention
+        mtsp.convert(model, [("dflash", config)])
+        return model
+
+    def test_sink_parameter_created_per_layer(self):
+        model = self._make_model()
+        heads = model.dflash_config.num_attention_heads
+        for layer in model.dflash_module.layers:
+            assert layer.self_attn.attention_sink_bias is not None
+            assert layer.self_attn.attention_sink_bias.shape == (heads,)
+            assert layer.self_attn.attention_sink_bias.requires_grad
+
+    def test_absent_by_default(self):
+        model = self._make_model(attention_sink=False)
+        for layer in model.dflash_module.layers:
+            assert layer.self_attn.attention_sink_bias is None
+
+    def test_sink_receives_gradient(self):
+        model = self._make_model()
+        torch.manual_seed(0)
+        input_ids = torch.randint(1, model.dflash_config.vocab_size, (2, SEQ_LEN))
+        model.train()
+        out = model(
+            input_ids=input_ids,
+            attention_mask=torch.ones_like(input_ids),
+            labels=input_ids.clone(),
+        )
+        assert torch.isfinite(out.loss)
+        out.loss.backward()
+        grads = [layer.self_attn.attention_sink_bias.grad for layer in model.dflash_module.layers]
+        assert all(g is not None and torch.isfinite(g).all() for g in grads)
+        assert any(g.abs().sum() > 0 for g in grads), "sink should receive a non-zero gradient"
+
+    def test_very_negative_sink_matches_no_sink(self):
+        """A sink at -inf carries no mass, so the layer must match plain attention.
+
+        Compared at the attention-layer level (not end-to-end loss) so the check isolates
+        the sink math from the rest of the pipeline.
+        """
+        torch.manual_seed(0)
+        model = self._make_model()
+        attn = model.dflash_module.layers[0].self_attn
+        heads, head_dim = attn.num_heads, attn.head_dim
+        kv_heads = attn.num_kv_heads
+
+        q = torch.randn(2, heads, BLOCK_SIZE, head_dim)
+        k = torch.randn(2, kv_heads, SEQ_LEN, head_dim)
+        v = torch.randn(2, kv_heads, SEQ_LEN, head_dim)
+
+        # A sink at -inf contributes no probability mass, so dropping its column leaves
+        # exactly the plain softmax attention distribution.
+        torch.nn.init.constant_(attn.attention_sink_bias, float("-inf"))
+        with torch.no_grad():
+            got = attn._sink_attention(q, k, v, None)
+
+            k_rep = repeat_kv(k, attn.num_key_value_groups)
+            v_rep = repeat_kv(v, attn.num_key_value_groups)
+            weights = torch.matmul(q, k_rep.transpose(2, 3)) * attn.scaling
+            expected = (
+                torch.matmul(torch.softmax(weights, dim=-1), v_rep).transpose(1, 2).contiguous()
+            )
+        assert torch.allclose(got, expected, atol=1e-6), (got - expected).abs().max()
+
+    def test_sink_absorbs_probability_mass(self):
+        """A finite sink strictly reduces the mass left for real tokens."""
+        torch.manual_seed(0)
+        model = self._make_model()
+        attn = model.dflash_module.layers[0].self_attn
+        q = torch.randn(1, attn.num_heads, BLOCK_SIZE, attn.head_dim)
+        k = torch.randn(1, attn.num_kv_heads, SEQ_LEN, attn.head_dim)
+        v = torch.randn(1, attn.num_kv_heads, SEQ_LEN, attn.head_dim)
+
+        outs = []
+        for value in (-8.0, 0.0, 8.0):
+            torch.nn.init.constant_(attn.attention_sink_bias, value)
+            with torch.no_grad():
+                outs.append(attn._sink_attention(q, k, v, None).abs().sum().item())
+        # More sink mass -> less mass on real tokens -> smaller output magnitude.
+        assert outs[0] > outs[1] > outs[2], outs
+
+    def test_export_includes_sink_weights_and_flag(self, tmp_path):
+        model = self._make_model()
+        export_dir = tmp_path / "exported"
+        model.get_exporter().export(export_dir)
+
+        sd = load_file(str(export_dir / "model.safetensors"))
+        heads = model.dflash_config.num_attention_heads
+        for i in range(model.dflash_config.num_hidden_layers):
+            key = f"layers.{i}.self_attn.attention_sink_bias"
+            assert key in sd, f"missing {key}"
+            assert sd[key].shape == (heads,)
+
+        with open(export_dir / "config.json") as f:
+            cfg = json.load(f)
+        assert cfg["dflash_config"]["attention_sink_bias"] is True
+        assert cfg["attention_sink_bias"] is True
+
+    def test_export_omits_sink_when_disabled(self, tmp_path):
+        model = self._make_model(attention_sink=False)
+        export_dir = tmp_path / "exported_nosink"
+        model.get_exporter().export(export_dir)
+        sd = load_file(str(export_dir / "model.safetensors"))
+        assert not any("attention_sink_bias" in k for k in sd)
+        with open(export_dir / "config.json") as f:
+            cfg = json.load(f)
+        assert "attention_sink_bias" not in cfg["dflash_config"]
+
+
+class TestMarkovHeadKeyRemap:
+    """Head weights load from either the flat or the nested `markov_head.` layout.
+
+    ModelOpt (and the upstream DeepSpec reference) keeps the head tensors flat on the
+    module and exports them that way. Some released drafters — notably
+    nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark — nest them under a
+    `markov_head.` parent, so loading accepts both.
+    """
+
+    def _module(self):
+        model = get_tiny_llama(num_hidden_layers=4)
+        mtsp.convert(model, [("dflash", _get_dspark_config())])
+        return model.dflash_module
+
+    def test_nested_keys_load(self):
+        module = self._module()
+        flat = module.state_dict()
+        nested = {
+            (
+                f"markov_head.{k}" if k.startswith(("markov_w1", "markov_w2")) else k
+            ): torch.full_like(v, 0.5) if k.startswith(("markov_w1", "markov_w2")) else v
+            for k, v in flat.items()
+        }
+        res = module.load_state_dict(nested, strict=True)
+        assert not res.missing_keys and not res.unexpected_keys
+        assert torch.allclose(
+            module.markov_w1.weight, torch.full_like(module.markov_w1.weight, 0.5)
+        )
+
+    def test_flat_keys_still_load(self):
+        """The exported (flat) layout keeps working unchanged."""
+        module = self._module()
+        flat = dict(module.state_dict())
+        flat["markov_w1.weight"] = torch.full_like(flat["markov_w1.weight"], 0.25)
+        res = module.load_state_dict(flat, strict=True)
+        assert not res.missing_keys and not res.unexpected_keys
+        assert torch.allclose(
+            module.markov_w1.weight, torch.full_like(module.markov_w1.weight, 0.25)
+        )
+
+    def test_flat_key_wins_over_nested(self):
+        """An explicit flat key is never clobbered by a nested duplicate."""
+        module = self._module()
+        sd = dict(module.state_dict())
+        sd["markov_w1.weight"] = torch.full_like(sd["markov_w1.weight"], 1.0)
+        sd["markov_head.markov_w1.weight"] = torch.full_like(sd["markov_w1.weight"], 9.0)
+        module.load_state_dict(sd, strict=True)
+        assert torch.allclose(
+            module.markov_w1.weight, torch.full_like(module.markov_w1.weight, 1.0)
+        )
+
+
+class TestInitCheckpoint:
+    """dflash_init_checkpoint warm-starts the draft module from exported weights."""
+
+    def _make_model(self, init_checkpoint=None, **overrides):
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = _get_dspark_config()
+        config.update(overrides)
+        if init_checkpoint is not None:
+            config["dflash_init_checkpoint"] = str(init_checkpoint)
+        mtsp.convert(model, [("dflash", config)])
+        return model
+
+    def _export(self, tmp_path, **overrides):
+        """Train-free export of a converted model, to be reloaded as a warm start."""
+        model = self._make_model(**overrides)
+        # Make the weights distinctive so a silent re-init would be visible.
+        with torch.no_grad():
+            for p in model.dflash_module.parameters():
+                p.fill_(0.125)
+        export_dir = tmp_path / "drafter"
+        model.get_exporter().export(export_dir)
+        return export_dir
+
+    def test_warm_start_loads_exported_weights(self, tmp_path):
+        """Every draft parameter comes from the checkpoint, not a fresh init."""
+        export_dir = self._export(tmp_path)
+        model = self._make_model(init_checkpoint=export_dir)
+        for name, p in model.dflash_module.named_parameters():
+            assert torch.allclose(p, torch.full_like(p, 0.125)), f"{name} was not warm-started"
+
+    def test_accepts_safetensors_file_path(self, tmp_path):
+        """The file itself works, not just its directory."""
+        export_dir = self._export(tmp_path)
+        model = self._make_model(init_checkpoint=export_dir / "model.safetensors")
+        assert torch.allclose(
+            model.dflash_module.fc.weight, torch.full_like(model.dflash_module.fc.weight, 0.125)
+        )
+
+    def test_round_trip_with_sink_and_causal(self, tmp_path):
+        """Warm start carries the sink weights of a causal + sink drafter."""
+        export_dir = self._export(
+            tmp_path, dflash_attention_sink=True, dflash_draft_attention="causal"
+        )
+        model = self._make_model(
+            init_checkpoint=export_dir,
+            dflash_attention_sink=True,
+            dflash_draft_attention="causal",
+        )
+        for layer in model.dflash_module.layers:
+            sink = layer.self_attn.attention_sink_bias
+            assert sink is not None
+            assert torch.allclose(sink, torch.full_like(sink, 0.125))
+
+    def test_default_is_random_init(self, tmp_path):
+        """Without the option nothing is loaded (regression guard for the default path)."""
+        self._export(tmp_path)
+        model = self._make_model()
+        assert not torch.allclose(
+            model.dflash_module.fc.weight, torch.full_like(model.dflash_module.fc.weight, 0.125)
+        )
+
+    def test_missing_path_raises(self, tmp_path):
+        with pytest.raises(FileNotFoundError, match="no draft weights"):
+            self._make_model(init_checkpoint=tmp_path / "does_not_exist")
+
+    def test_architecture_mismatch_raises(self, tmp_path):
+        """A checkpoint for a different draft depth must not partially load.
+
+        Depth feeds `fc`'s input width (one target hidden per draft layer), so this trips
+        the shape check; either failure mode is acceptable as long as it raises.
+        """
+        export_dir = self._export(tmp_path)
+        with pytest.raises(ValueError, match=r"shape mismatch|does not match the configured draft"):
+            self._make_model(
+                init_checkpoint=export_dir,
+                dflash_architecture_config={
+                    **_get_dspark_config()["dflash_architecture_config"],
+                    "num_hidden_layers": NUM_DRAFT_LAYERS + 1,
+                },
+            )
+
+    def test_sink_mismatch_raises(self, tmp_path):
+        """Exported-without-sink cannot warm-start a sink-enabled draft."""
+        export_dir = self._export(tmp_path, dflash_attention_sink=False)
+        with pytest.raises(ValueError, match="does not match the configured draft"):
+            self._make_model(init_checkpoint=export_dir, dflash_attention_sink=True)
+
+    def test_nested_head_shape_mismatch_reported(self, tmp_path):
+        """A wrong-shaped tensor is caught even under the nested `markov_head.` layout.
+
+        The shape check has to resolve the module's load hooks first; otherwise a remapped
+        key skips it and fails later with a far less obvious error.
+        """
+        export_dir = self._export(tmp_path)
+        path = export_dir / "model.safetensors"
+        sd = load_file(str(path))
+        sd["markov_head.markov_w1.weight"] = torch.zeros(3, MARKOV_RANK)
+        del sd["markov_w1.weight"]
+        save_file(sd, str(path))
+        with pytest.raises(ValueError, match="shape mismatch"):
+            self._make_model(init_checkpoint=export_dir)
+
+
+class TestExplicitTargetLayerIds:
+    """dflash_architecture_config.target_layer_ids overrides the uniform default.
+
+    A published draft is trained against specific capture points; recomputing the default
+    would feed it features from layers it never saw (and mis-shape ``fc``).
+    """
+
+    def _make_model(self, target_layer_ids=None, num_layers=NUM_DRAFT_LAYERS):
+        model = get_tiny_llama(num_hidden_layers=8)
+        config = _get_dspark_config(num_layers=num_layers)
+        if target_layer_ids is not None:
+            config["dflash_architecture_config"]["target_layer_ids"] = target_layer_ids
+        mtsp.convert(model, [("dflash", config)])
+        return model
+
+    def test_explicit_ids_are_used(self):
+        model = self._make_model(target_layer_ids=[0, 7])
+        assert model.target_layer_ids == [0, 7]
+        assert model.dflash_config.target_layer_ids == [0, 7]
+
+    def test_default_when_unset(self):
+        """Without an override the uniform default is still derived."""
+        model = self._make_model()
+        assert len(model.target_layer_ids) == NUM_DRAFT_LAYERS
+        assert model.target_layer_ids == build_target_layer_ids(8, NUM_DRAFT_LAYERS)
+
+    def test_wrong_count_raises(self):
+        with pytest.raises(ValueError, match="one target layer per draft layer"):
+            self._make_model(target_layer_ids=[0, 3, 7])
+
+    def test_out_of_range_raises(self):
+        with pytest.raises(ValueError, match="beyond the base model"):
+            self._make_model(target_layer_ids=[0, 99])
+
+    def test_export_round_trips_explicit_ids(self, tmp_path):
+        model = self._make_model(target_layer_ids=[0, 7])
+        model.get_exporter().export(tmp_path / "exp")
+        with open(tmp_path / "exp" / "config.json") as f:
+            cfg = json.load(f)
+        assert cfg["dflash_config"]["target_layer_ids"] == [0, 7]