[Speculative Decoding] DFlash2 draft variant (sublayer convolution + candidate selector)

.pre-commit-config.yaml

@@ -116,6 +116,8 @@ repos:
               modelopt/torch/speculative/plugins/modeling_domino.py|
               modelopt/torch/speculative/plugins/hf_dflash.py|
               modelopt/torch/speculative/plugins/modeling_dflash.py|
+              modelopt/torch/speculative/plugins/hf_dflash2.py|
+              modelopt/torch/speculative/plugins/modeling_dflash2.py|
               modelopt/torch/speculative/plugins/hf_dspark.py|
               modelopt/torch/speculative/plugins/modeling_dspark.py|
               modelopt/torch/speculative/plugins/hf_medusa.py|

CHANGELOG.rst

@@ -16,6 +16,10 @@ Changelog
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
 
+*Speculative Decoding*
+
+- Add the DFlash2 draft variant, selected with ``dflash_architecture_config.projector_type="dflash2"``. It keeps DFlash's one-pass parallel backbone and adds a grouped dynamic convolution around every attention/MLP sublayer (sized by ``conv_kernel_size`` / ``conv_group_size``) plus a low-rank candidate selector (``selector_rank`` / ``selector_top_k``) that scores transitions between adjacent block positions' top-k candidates. The selector's training term is weighted by ``dflash_selector_loss_alpha`` (default 1.0). Exported checkpoints declare ``DFlash2DraftModel`` and match the SGLang/vLLM DFlash2 loaders.
+
 *Misc*
 
 - Add ``modelopt.torch.utils.mlflow.MlflowRunLogger`` for recording a script run on an MLflow tracking server: the invocation, the ModelOpt version, the run log (captured by teeing ``stdout``/``stderr``) and any caller-supplied artifacts, with configuration as searchable params. ``mlflow`` is an optional dependency, imported only when tracking is enabled.

modelopt/torch/export/plugins/hf_spec_export.py

@@ -31,18 +31,26 @@
 
 
 def _get_rope_theta(config, default=None):
-    """Get RoPE theta from either legacy or Transformers 5 config fields."""
-    rope_theta = getattr(config, "rope_theta", None)
-    if rope_theta is not None:
-        return rope_theta
-
+    """Get RoPE theta from either legacy or Transformers 5 config fields.
+
+    ``rope_parameters`` is checked FIRST. A config can carry both fields with
+    different values: Transformers 5 stores the real base under
+    ``rope_parameters`` while the class default (10000.0 for Qwen3) may still be
+    visible as a top-level ``rope_theta``. Reading ``rope_theta`` first silently
+    exports a draft whose RoPE base is 100x off the target's, which breaks
+    serving because DFlash injects the target's KV into every draft layer.
+    """
     # Transformers 5 stores this under rope_parameters (and exposes the same
     # data through rope_scaling for backwards compatibility).
     for attr in ("rope_parameters", "rope_scaling"):
         rope_config = getattr(config, attr, None)
         if isinstance(rope_config, dict) and rope_config.get("rope_theta") is not None:
             return rope_config["rope_theta"]
 
+    rope_theta = getattr(config, "rope_theta", None)
+    if rope_theta is not None:
+        return rope_theta
+
     return default
 
 
@@ -533,3 +541,45 @@ def _export_config(self):
             }
         )
         return config
+
+
+class DFlash2Exporter(DFlashExporter):
+    """Draft model exporter for DFlash2 (DFlash backbone + convolutions + selector).
+
+    Same z-lab-compatible format as DFlash, plus the DFlash2 weights
+    (``layers.*.attention_conv.*`` / ``layers.*.mlp_conv.*`` /
+    ``candidate_selector.*``, already captured by the inherited ``dflash_module.``
+    stripping) and the config fields the SGLang/vLLM ``DFlash2DraftModel`` loader
+    needs to rebuild them (``conv_kernel_size``, ``conv_group_size``,
+    ``selector_rank``, ``selector_top_k``).
+
+    The architecture name is what selects the DFlash2 serving path: a checkpoint
+    declaring ``DFlashDraftModel`` loads as a plain DFlash draft and would silently
+    ignore the convolutions and the selector.
+    """
+
+    def _export_config(self):
+        """Extend the DFlash config with the DFlash2 architecture fields."""
+        config = super()._export_config()
+        draft_config = self.model.dflash_config
+
+        config["architectures"] = ["DFlash2DraftModel"]
+        # Present because HFDFlash2Model.modify validates them at convert time.
+        config["dflash_config"].update(
+            {
+                "projector_type": getattr(draft_config, "projector_type", "dflash2"),
+                "conv_kernel_size": draft_config.conv_kernel_size,
+                "conv_group_size": draft_config.conv_group_size,
+                "selector_rank": draft_config.selector_rank,
+                "selector_top_k": draft_config.selector_top_k,
+                # The published DFlash2 checkpoints carry block_size inside
+                # dflash_config; the DFlash loader reads it from the top level.
+                # Emit both so either contract resolves to the same value.
+                "block_size": config["block_size"],
+            }
+        )
+        # Published DFlash2 checkpoints state causality explicitly rather than
+        # leaving it to be inferred from layer_types. Only set it when the SWA
+        # block above has not already written a `causal` entry.
+        config.setdefault("is_causal", config["dflash_config"].get("causal", False))
+        return config

modelopt/torch/speculative/config.py

@@ -228,6 +228,17 @@ class DFlashConfig(ModeloptBaseConfig):
         ),
     )
 
+    dflash_selector_loss_alpha: float = ModeloptField(
+        default=1.0,
+        ge=0.0,
+        description=(
+            "DFlash2 only: weight of the candidate-selector cross-entropy term, added to "
+            "the backbone loss. The selector re-ranks the backbone's top-k candidates per "
+            "block position; 0 trains the backbone and convolutions only. "
+            "Ignored unless dflash_architecture_config.projector_type == 'dflash2'."
+        ),
+    )
+
     @model_validator(mode="after")
     def _check_dpace_alpha(self) -> "DFlashConfig":
         # Validate at construction regardless of the active objective, so a bad alpha

modelopt/torch/speculative/dflash/conversion.py

@@ -35,6 +35,12 @@
 # ``dflash_architecture_config.projector_type == "dspark"`` and kept in its own
 # registry so its wrapper (HFDSparkModel) does not overwrite HFDFlashModel.
 DSparkDMRegistry = _DMRegistryCls(prefix="DSpark")
+# DFlash2 also reuses the dflash mode/config/recipe, converting the base model to a
+# DFlash backbone whose sublayers are wrapped in grouped dynamic convolutions, plus a
+# low-rank candidate selector. Selected via
+# ``dflash_architecture_config.projector_type == "dflash2"`` and kept in its own
+# registry so its wrapper (HFDFlash2Model) does not overwrite HFDFlashModel.
+DFlash2DMRegistry = _DMRegistryCls(prefix="DFlash2")
 
 
 def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertReturnType:
@@ -53,12 +59,14 @@ def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertRe
         registry = DominoDMRegistry
     elif projector_type == "dspark":
         registry = DSparkDMRegistry
+    elif projector_type == "dflash2":
+        registry = DFlash2DMRegistry
     elif projector_type in (None, "dflash"):
         registry = DFlashDMRegistry
     else:
         raise ValueError(
             f"Unsupported dflash_architecture_config.projector_type: {projector_type!r}. "
-            "Expected 'dflash' (default), 'domino' or 'dspark'."
+            "Expected 'dflash' (default), 'domino', 'dspark' or 'dflash2'."
         )
 
     original_cls = type(model)

modelopt/torch/speculative/plugins/__init__.py

@@ -31,6 +31,7 @@
 
 with import_plugin("transformers"):
     from .hf_dflash import *
+    from .hf_dflash2 import *
     from .hf_domino import *
     from .hf_dspark import *
     from .hf_eagle import *

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -419,10 +419,21 @@ def modify(self, config):
         # overwrite any user value and warn. (rope_scaling is intentionally NOT inherited:
         # DFlash uses standard Qwen3 RotaryEmbedding; the long-context YaRN scaling is
         # added only at export via dflash_export_rope_scaling.)
+        # A config can carry BOTH a top-level rope_theta and a rope_parameters dict
+        # with different values: Transformers 5 keeps the real base in
+        # rope_parameters while the class default (10000.0 for Qwen3) stays visible
+        # as rope_theta. rope_parameters wins, otherwise the draft trains against a
+        # RoPE base 100x off the target's.
+        base_rope_params = getattr(base_config, "rope_parameters", None)
+        if not isinstance(base_rope_params, dict):
+            base_rope_params = {}
         for attr in ("rope_theta", "rope_type", "rope_interleaved"):
-            if not hasattr(base_config, attr):
+            if attr in base_rope_params:
+                base_val = base_rope_params[attr]
+            elif hasattr(base_config, attr):
+                base_val = getattr(base_config, attr)
+            else:
                 continue
-            base_val = getattr(base_config, attr)
             user_val = getattr(self.dflash_config, attr, None)
             if user_val is not None and user_val != base_val:
                 logger.warning(
@@ -434,6 +445,12 @@ def modify(self, config):
                     base_val,
                 )
             setattr(self.dflash_config, attr, base_val)
+            # Qwen3Config populates rope_parameters at construction, so a later
+            # setattr on the flat field alone would leave the dict — which is what
+            # the rotary module reads — holding the stale value.
+            draft_rope_params = getattr(self.dflash_config, "rope_parameters", None)
+            if isinstance(draft_rope_params, dict) and attr in draft_rope_params:
+                draft_rope_params[attr] = base_val
 
         self.dflash_config.head_dim = getattr(
             self.dflash_config,
@@ -632,7 +649,14 @@ def _build_generate_swa_mask(self, ctx_len, bsz, dtype, device):
         return attn_mask
 
     def _compute_loss(
-        self, logits, input_ids, anchor_positions, block_keep_mask, loss_mask, base_logits=None
+        self,
+        logits,
+        input_ids,
+        anchor_positions,
+        block_keep_mask,
+        loss_mask,
+        base_logits=None,
+        draft_hidden=None,
     ):
         """Compute weighted cross-entropy (or KD) loss and accuracy.
 
@@ -643,6 +667,8 @@ def _compute_loss(
             block_keep_mask: Valid block mask [B, N].
             loss_mask: Token-level loss mask [B, seq_len].
             base_logits: Base model logits for KD loss [B, seq_len, vocab], or None for CE.
+            draft_hidden: Draft hidden states [B, N*block_size, H] behind ``logits``.
+                Unused here; DFlash2 needs them for its candidate-selector term.
 
         Returns:
             (loss, accuracy) tuple.
@@ -921,6 +947,7 @@ def forward(
             block_keep_mask,
             loss_mask,
             base_outputs.logits if self.dflash_self_logit_distillation else None,
+            draft_hidden=hidden,
         )
 
         return ModelOutput(