Support fine-tuning released DFlash/DSpark drafters (causal SWA, attention sink, warm start)

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,

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

modelopt/torch/export/plugins/hf_spec_export.py

@@ -412,21 +412,34 @@ 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
             config["dflash_config"].update(
                 {
                     "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)

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 >= "

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