From 1419d47e5a775bddaec7f08c5eeff34864abc68c Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 19 Aug 2026 08:41:05 +0000 Subject: [PATCH 01/10] refactor(speculative): add no-op sublayer wrappers to the DFlash draft layer DFlash2 wraps every attention and MLP sublayer in a grouped dynamic convolution. Give DFlashDecoderLayer a prepare()/finish() seam around each sublayer so a variant can transform the sublayer's input and output without the layer's forward growing a branch. The default wrapper is a parameterless no-op, so DFlash, Domino and DSpark drafts keep their exact numerics and state_dict contents. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../speculative/plugins/modeling_dflash.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py index 6463cb4109d..6bcff38d6c6 100644 --- a/modelopt/torch/speculative/plugins/modeling_dflash.py +++ b/modelopt/torch/speculative/plugins/modeling_dflash.py @@ -221,6 +221,27 @@ def forward(self, hidden_states, target_hidden, position_embeddings, attention_m return self.o_proj(attn_output) +class _IdentitySublayerWrapper(nn.Module): + """No-op sublayer wrapper: the default around each attention/MLP sublayer. + + ``DFlashDecoderLayer`` calls ``prepare()`` before a sublayer and ``finish()`` + after it, so a variant can transform the sublayer's input and output without + the layer's forward growing a branch. This default does nothing and holds no + parameters, so it neither appears in ``state_dict()`` nor changes the numerics + of a plain DFlash (or Domino/DSpark) draft. + + DFlash2 substitutes ``DFlashGroupedConv`` here (see ``modeling_dflash2.py``). + """ + + def prepare(self, hidden_states): + """Return the sublayer input unchanged, with no state to carry to ``finish``.""" + return hidden_states, None + + def finish(self, hidden_states, state): + """Return the sublayer output unchanged.""" + return hidden_states + + class DFlashDecoderLayer(nn.Module): """Draft decoder layer with KV injection.""" @@ -231,19 +252,26 @@ def __init__(self, config, layer_idx): self.mlp = _MLP_CLS(config) self.input_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps) + # Sublayer wrappers; no-ops unless a variant replaces them (DFlash2). + self.attention_conv = _IdentitySublayerWrapper() + self.mlp_conv = _IdentitySublayerWrapper() def forward(self, hidden_states, target_hidden, position_embeddings, attention_mask=None): """Forward pass with residual connections.""" residual = hidden_states hidden_states = self.input_layernorm(hidden_states) + hidden_states, conv_state = self.attention_conv.prepare(hidden_states) hidden_states = self.self_attn( hidden_states, target_hidden, position_embeddings, attention_mask ) + hidden_states = self.attention_conv.finish(hidden_states, conv_state) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, conv_state = self.mlp_conv.prepare(hidden_states) hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp_conv.finish(hidden_states, conv_state) hidden_states = residual + hidden_states return hidden_states From 6b4ca0d1a9999f48d1befa74a8dc1db8b1b40d88 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 19 Aug 2026 08:52:11 +0000 Subject: [PATCH 02/10] feat(speculative): DFlash2 draft variant (sublayer conv + candidate selector) DFlash2 (https://inco.ai/blog/dflash2/) keeps DFlash's one-pass parallel backbone and adds two components that recover the acceptance a purely parallel draft loses: - a grouped dynamic depthwise convolution around every attention and MLP sublayer, giving each block position a view of its predecessors inside the block (taps do not cross the block boundary); - a low-rank candidate selector that scores transitions between adjacent positions' top-k candidates, so serving walks one coherent path instead of taking an independent argmax per position. Selected with dflash_architecture_config.projector_type='dflash2', alongside 'domino' and 'dspark'. The convolution's base kernel is identity-initialized, so a fresh DFlash2 draft starts out computing exactly what its DFlash backbone would. The selector is supervised by a cross-entropy term over its candidate set, weighted by dflash_selector_loss_alpha. Positions are scored against their teacher-forced predecessor so they train in parallel, and the gold token is substituted into the candidate set when the backbone's top-k misses it. Module and parameter names match the SGLang/vLLM DFlash2DraftModel loaders, so an exported checkpoint is served directly. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../torch/export/plugins/hf_spec_export.py | 34 ++ modelopt/torch/speculative/config.py | 11 + .../torch/speculative/dflash/conversion.py | 10 +- .../torch/speculative/plugins/__init__.py | 1 + .../torch/speculative/plugins/hf_dflash.py | 12 +- .../torch/speculative/plugins/hf_dflash2.py | 198 +++++++++++ .../speculative/plugins/modeling_dflash2.py | 314 ++++++++++++++++++ 7 files changed, 578 insertions(+), 2 deletions(-) create mode 100644 modelopt/torch/speculative/plugins/hf_dflash2.py create mode 100644 modelopt/torch/speculative/plugins/modeling_dflash2.py diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 255b1d9ab04..820c5d58f53 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -533,3 +533,37 @@ 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, + } + ) + return config diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index 4df57d3035f..bd199652392 100644 --- a/modelopt/torch/speculative/config.py +++ b/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 diff --git a/modelopt/torch/speculative/dflash/conversion.py b/modelopt/torch/speculative/dflash/conversion.py index 0516ab7181d..a2fc92bd0d5 100644 --- a/modelopt/torch/speculative/dflash/conversion.py +++ b/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) diff --git a/modelopt/torch/speculative/plugins/__init__.py b/modelopt/torch/speculative/plugins/__init__.py index ea789388ad4..87f66947bd2 100644 --- a/modelopt/torch/speculative/plugins/__init__.py +++ b/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 * diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index e0d63bde136..fe13f9c6b2a 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -632,7 +632,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 +650,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 +930,7 @@ def forward( block_keep_mask, loss_mask, base_outputs.logits if self.dflash_self_logit_distillation else None, + draft_hidden=hidden, ) return ModelOutput( diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py new file mode 100644 index 00000000000..fb4ede5c70a --- /dev/null +++ b/modelopt/torch/speculative/plugins/hf_dflash2.py @@ -0,0 +1,198 @@ +# 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. + +"""HF DFlash2 model wrapper — DFlash training plus the candidate-selector objective. + +DFlash2 differs from DFlash only in the draft module (grouped dynamic convolutions +around every sublayer, plus a candidate selector) and in one extra loss term, so +this wrapper reuses ``HFDFlashModel``'s forward wholesale and overrides just +:meth:`_compute_loss`. + +The convolutions need no supervision of their own: they sit inside the backbone +and are trained by the backbone loss. The selector does, because at serving time +it — not an independent argmax — picks the drafted token at each block position. + +Selector supervision (following the SGLang/SpecForge reference): + +- Take the backbone's top-k candidates per block position. +- Score each candidate against its *teacher-forced* predecessor token, so the + positions train in parallel exactly as the backbone does. +- When the gold token is missing from the top-k, substitute it into the last + candidate slot. Without this the selector sees no positive class on the hard + positions and never learns those edges. +""" + +import torch +import torch.nn.functional as F +from transformers import PreTrainedModel + +from ..dflash.conversion import DFlash2DMRegistry +from .hf_dflash import HFDFlashModel +from .modeling_dflash2 import DFlash2Module + +__all__ = ["HFDFlash2Model"] + + +@DFlash2DMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) +class HFDFlash2Model(HFDFlashModel): + """DFlash model with DFlash2's sublayer convolutions and candidate selector. + + Registered in ``DFlash2DMRegistry`` so that ``convert_to_dflash_model`` routes + to it when ``dflash_architecture_config.projector_type == "dflash2"``. + """ + + def _build_draft_module(self, dflash_config): + """Build the DFlash2 draft module (DFlash backbone + convolutions + selector).""" + return DFlash2Module(dflash_config) + + def modify(self, config): + """Initialize the DFlash2 draft module and read the selector loss weight.""" + arch_config = config.dflash_architecture_config + missing = [ + name + for name in ("conv_kernel_size", "conv_group_size", "selector_rank", "selector_top_k") + if arch_config.get(name) is None + ] + if missing: + raise ValueError( + f"DFlash2 (projector_type='dflash2') requires {missing} in " + "dflash_architecture_config (convolution taps/group size and the " + "candidate selector's rank/top-k)." + ) + super().modify(config) + self.dflash_selector_loss_alpha = getattr(config, "dflash_selector_loss_alpha", 1.0) + + def get_exporter(self): + """Get the exporter for the DFlash2 draft model.""" + from modelopt.torch.export.plugins.hf_spec_export import DFlash2Exporter + + return DFlash2Exporter(self) + + def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mask): + """Cross-entropy over the selector's candidate set, and its top-1 accuracy. + + Args: + logits: Backbone logits per block position ``[B, N, block_size, V]``. + target_ids: Gold token ids ``[B, N, block_size]``. + hidden: Backbone hidden states ``[B, N, block_size, H]``. + predecessor_ids: Teacher-forced predecessor ids ``[B, N, block_size]``. + weight_mask: Per-position loss weights ``[B, N, block_size]``. + + Returns: + ``(loss, accuracy, coverage)`` — coverage is the fraction of supervised + positions whose gold token was already in the backbone's top-k, i.e. how + often the selector is choosing rather than being handed the answer. + """ + selector = self.dflash_module.candidate_selector + top_k = selector.top_k + + unary_logits, candidate_ids = logits.topk(top_k, dim=-1) + + # Where the gold token is absent from the top-k, overwrite the last (lowest + # scoring) slot with it, so every supervised position has a correct class. + gold_in_topk = (candidate_ids == target_ids.unsqueeze(-1)).any(dim=-1) + gold_slot = torch.where( + gold_in_topk, + (candidate_ids == target_ids.unsqueeze(-1)).float().argmax(dim=-1), + torch.full_like(target_ids, top_k - 1), + ) + gold_unary = logits.gather(-1, target_ids.unsqueeze(-1)) + candidate_ids = candidate_ids.scatter(-1, gold_slot.unsqueeze(-1), target_ids.unsqueeze(-1)) + unary_logits = unary_logits.scatter(-1, gold_slot.unsqueeze(-1), gold_unary) + + selector_logits = selector.score_candidates( + candidate_ids, unary_logits, hidden, predecessor_ids + ) + + flat_weights = weight_mask.reshape(-1) + denominator = flat_weights.sum() + 1e-6 + per_token = F.cross_entropy( + selector_logits.float().reshape(-1, top_k), + gold_slot.reshape(-1), + reduction="none", + ) + loss = (per_token * flat_weights).sum() / denominator + + with torch.no_grad(): + chosen = selector_logits.argmax(dim=-1).reshape(-1) + accuracy = ( + (chosen == gold_slot.reshape(-1)).float() * flat_weights + ).sum() / denominator + coverage = (gold_in_topk.reshape(-1).float() * flat_weights).sum() / denominator + return loss, accuracy.item(), coverage.item() + + def _compute_loss( + self, + logits, + input_ids, + anchor_positions, + block_keep_mask, + loss_mask, + base_logits=None, + draft_hidden=None, + ): + """Backbone DFlash loss plus the candidate-selector cross-entropy. + + Reuses ``HFDFlashModel._compute_loss`` for the backbone term, then rebuilds + the same target/weight alignment for the selector term. Reported accuracy + stays the backbone's top-1, so DFlash and DFlash2 runs remain comparable; + the selector's own accuracy is logged separately. + """ + loss, accuracy = super()._compute_loss( + logits, input_ids, anchor_positions, block_keep_mask, loss_mask, base_logits + ) + if self.dflash_selector_loss_alpha <= 0 or draft_hidden is None: + return loss, accuracy + + bsz, seq_len = input_ids.shape + block_size = self.dflash_block_size + n_blocks = anchor_positions.shape[1] + device = input_ids.device + + offsets = torch.arange(block_size, device=device).view(1, 1, -1) + label_indices = anchor_positions.unsqueeze(-1) + offsets + valid_label = label_indices < seq_len + safe_label_indices = label_indices.clamp(max=seq_len - 1) + expanded_ids = input_ids.unsqueeze(1).expand(-1, n_blocks, -1) + target_ids = torch.gather(expanded_ids, 2, safe_label_indices) + + # Same supervision mask as the backbone loss: valid block, in bounds, not the + # anchor slot, and inside the answer span. Position weighting (decay/D-PACE) is + # deliberately not applied — it shapes *where* the backbone spends capacity, + # while the selector should learn every position's transition equally. + weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float() + weight_mask = weight_mask * valid_label.float() + weight_mask = weight_mask * (offsets > 0).float() + weight_mask = weight_mask * torch.gather( + loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices + ) + + # Teacher-forced predecessor of block position k is the real token at anchor+k-1; + # position 0's predecessor is the anchor itself, matching the serving-side walk + # which starts from the last verified token. + predecessor_ids = torch.gather(expanded_ids, 2, (safe_label_indices - 1).clamp(min=0)) + + selector_loss, selector_accuracy, selector_coverage = self._selector_loss( + logits.reshape(bsz, n_blocks, block_size, -1), + target_ids, + draft_hidden.reshape(bsz, n_blocks, block_size, -1), + predecessor_ids, + weight_mask, + ) + self._selector_metrics = { + "selector_accuracy": selector_accuracy, + "selector_coverage": selector_coverage, + } + return loss + self.dflash_selector_loss_alpha * selector_loss, accuracy diff --git a/modelopt/torch/speculative/plugins/modeling_dflash2.py b/modelopt/torch/speculative/plugins/modeling_dflash2.py new file mode 100644 index 00000000000..309bb4c54e2 --- /dev/null +++ b/modelopt/torch/speculative/plugins/modeling_dflash2.py @@ -0,0 +1,314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +# Adapted from https://github.com/sgl-project/SpecForge/pull/772 +# Copyright (c) 2025 sgl-project +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# 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. + +"""DFlash2 draft module — DFlash backbone plus local convolution and candidate selection. + +DFlash2 (Inco AI / Z Lab, https://inco.ai/blog/dflash2/) keeps DFlash's one-pass +parallel backbone and adds two small components that address the two ways a +purely parallel draft loses acceptance: + +- :class:`DFlashGroupedConv` — a grouped *dynamic* depthwise convolution wrapped + around every attention and MLP sublayer. Each block position mixes in its + predecessors inside the block, which injects the intra-block sequential + dependency the parallel backbone lacks (mitigating suffix acceptance decay) + without a second backbone pass. Taps do not cross the block boundary. + +- :class:`CandidateSelector` — a low-rank transition scorer. Instead of an + independent argmax per block position, the drafter keeps the target head's + top-k candidates per position and scores adjacent transitions, so serving can + walk one coherent path through the block. + +Where Domino uses a GRU and DSpark a Markov transition bias, DFlash2 spends its +extra capacity on these two pieces: both are cheap (a few percent of draft +parameters, ~1% of serving step latency in the reference measurements). + +This module owns the parameters only; the training wrapper (``HFDFlash2Model`` +in ``hf_dflash2.py``) orchestrates the forward and the selector loss. Module and +parameter names (``attention_conv`` / ``mlp_conv`` / ``base_kernel`` / +``kernel_projection`` / ``candidate_selector`` / ``predecessor_codebook`` / +``successor_codebook`` / ``hidden_projection``) match the SGLang and vLLM +``DFlash2DraftModel`` loaders so an exported checkpoint is served directly. +""" + +import torch +import torch.nn.functional as F +from torch import nn + +from .modeling_dflash import DFlashModule + +__all__ = ["CandidateSelector", "DFlash2Module", "DFlashGroupedConv"] + + +class DFlashGroupedConv(nn.Module): + """Grouped dynamic depthwise convolution over positions within a proposal block. + + Wraps one sublayer: :meth:`prepare` convolves the sublayer input and emits the + dynamic kernel for the output side, :meth:`finish` convolves the sublayer + output. One projection of the sublayer input produces both sides' kernel + deltas. + + ``base_kernel`` starts as an identity (tap 0 weight 1, later taps 0), so a + freshly built DFlash2 draft computes exactly what its DFlash backbone would. + That makes the convolution a stable extension rather than a perturbation, and + lets a DFlash checkpoint warm-start a DFlash2 run. + """ + + def __init__(self, hidden_size: int, block_size: int, taps: int, group_size: int): + """Build the identity-initialized base kernel and the dynamic-kernel projection.""" + super().__init__() + if taps < 1: + raise ValueError(f"DFlash2 conv_kernel_size must be >= 1, got {taps}.") + if taps > block_size: + raise ValueError( + f"DFlash2 conv_kernel_size ({taps}) must not exceed " + f"dflash_block_size ({block_size})." + ) + if group_size < 1 or hidden_size % group_size: + raise ValueError( + f"DFlash2 conv_group_size ({group_size}) must be >= 1 and divide " + f"hidden_size ({hidden_size})." + ) + + self.block_size = int(block_size) + self.taps = int(taps) + self.group_size = int(group_size) + self.num_groups = int(hidden_size) // self.group_size + + # [input/output side, tap, channel]; identity at tap 0. Layout matches the + # SGLang/vLLM DFlash2 weight loader. + base_kernel = torch.zeros(2, self.taps, int(hidden_size)) + base_kernel[:, 0] = 1.0 + self.base_kernel = nn.Parameter(base_kernel) + self.kernel_projection = nn.Linear( + int(hidden_size), 2 * self.taps * self.num_groups, bias=False + ) + + def _convolve(self, hidden_states, delta, side: int): + """Apply the depthwise convolution for one side, with taps clipped at block starts.""" + bsz, seq_len, hidden_size = hidden_states.shape + if seq_len % self.block_size: + raise ValueError( + f"DFlash2 convolution needs a sequence length divisible by " + f"block_size ({self.block_size}), got {seq_len}." + ) + + n_blocks = seq_len // self.block_size + blocks = hidden_states.reshape( + bsz, n_blocks, self.block_size, self.num_groups, self.group_size + ) + dynamic = delta.reshape(bsz, n_blocks, self.block_size, self.taps, self.num_groups) + base = self.base_kernel[side].reshape(1, 1, 1, self.taps, self.num_groups, self.group_size) + # Per-position, per-group coefficients: static base plus the dynamic delta. + coefficients = base + dynamic.unsqueeze(-1) + + output = coefficients[:, :, :, 0] * blocks + for tap in range(1, self.taps): + # Shift within the block only: position k reads k-tap, and the first + # `tap` positions of each block read zeros rather than the previous block. + shifted = F.pad(blocks[:, :, : self.block_size - tap], (0, 0, 0, 0, tap, 0)) + output = output + coefficients[:, :, :, tap] * shifted + return output.reshape(bsz, seq_len, hidden_size) + + def prepare(self, hidden_states): + """Convolve the sublayer input; return it with the output side's dynamic kernel.""" + coefficients = self.kernel_projection(hidden_states).reshape( + *hidden_states.shape[:-1], 2, self.taps, self.num_groups + ) + return self._convolve(hidden_states, coefficients[..., 0, :, :], side=0), coefficients[ + ..., 1, :, : + ] + + def finish(self, hidden_states, state): + """Convolve the sublayer output using the kernel produced by :meth:`prepare`.""" + return self._convolve(hidden_states, state, side=1) + + +class CandidateSelector(nn.Module): + """Low-rank scorer for transitions between adjacent block positions' candidates. + + Scores an edge from a predecessor token ``p`` to a candidate token ``c`` at a + block position with hidden state ``h`` as:: + + edge(p -> c) = + unary_logit[c] + + i.e. a bilinear form between the two token codebooks, gated by the context. + Training scores each position's candidate set independently under teacher + forcing (:meth:`score_candidates`); serving walks the resulting lattice. + """ + + def __init__(self, hidden_size: int, vocab_size: int, rank: int, top_k: int, std: float): + """Build the predecessor/successor codebooks and the context projection.""" + super().__init__() + if rank < 1: + raise ValueError(f"DFlash2 selector_rank must be >= 1, got {rank}.") + if not 1 <= top_k <= vocab_size: + raise ValueError( + f"DFlash2 selector_top_k must be in [1, vocab_size={vocab_size}], got {top_k}." + ) + self.top_k = int(top_k) + self.rank = int(rank) + self.predecessor_codebook = nn.Parameter(torch.empty(int(vocab_size), int(rank))) + self.successor_codebook = nn.Parameter(torch.empty(int(vocab_size), int(rank))) + self.hidden_projection = nn.Linear(int(hidden_size), int(rank), bias=False) + nn.init.normal_(self.predecessor_codebook, std=std) + nn.init.normal_(self.successor_codebook, std=std) + + def score_candidates(self, candidate_ids, unary_logits, hidden_states, predecessor_ids): + """Add the predecessor transition score to a candidate set's unary logits. + + Args: + candidate_ids: Candidate token ids ``[..., K]``. + unary_logits: Backbone logits for those candidates ``[..., K]``. + hidden_states: Backbone hidden at this position ``[..., H]``. + predecessor_ids: Teacher-forced predecessor token ids ``[...]``. + + Returns: + Selector logits over the candidate set ``[..., K]``. + """ + predecessor = self.predecessor_codebook[predecessor_ids] + successor = self.successor_codebook[candidate_ids] + context = predecessor * self.hidden_projection(hidden_states) + transition = torch.einsum("...r,...kr->...k", context.to(successor.dtype), successor) + return unary_logits + transition + + @torch.no_grad() + def greedy_path(self, candidate_ids, unary_logits, hidden_states, anchor_token_ids): + """Walk the candidate lattice greedily, mirroring the serving-side path walk. + + Args: + candidate_ids: ``[B, L, K]`` candidate ids per block position. + unary_logits: ``[B, L, K]`` backbone logits for those candidates. + hidden_states: ``[B, L, H]`` backbone hidden per block position. + anchor_token_ids: ``[B]`` the verified token preceding the block. + + Returns: + Selected token ids ``[B, L]``. + """ + predecessor_ids = anchor_token_ids + path = [] + for position in range(candidate_ids.shape[1]): + scores = self.score_candidates( + candidate_ids[:, position], + unary_logits[:, position], + hidden_states[:, position], + predecessor_ids, + ) + selected = scores.argmax(dim=-1, keepdim=True) + predecessor_ids = candidate_ids[:, position].gather(1, selected)[:, 0] + path.append(predecessor_ids) + return torch.stack(path, dim=1) + + +class DFlash2Module(DFlashModule): + """DFlash draft backbone with per-sublayer convolutions and a candidate selector.""" + + def __init__(self, config): + """Initialize the DFlash backbone, then attach the convolutions and the selector.""" + super().__init__(config) + + self.projector_type = getattr(config, "projector_type", "dflash2") + + def required_int(name: str) -> int: + """Read an int architecture field, rejecting missing values and bools.""" + value = getattr(config, name, None) + if not isinstance(value, int) or isinstance(value, bool): + raise ValueError( + f"DFlash2 (projector_type='dflash2') requires an integer " + f"'{name}' in dflash_architecture_config, got {value!r}." + ) + return value + + taps = required_int("conv_kernel_size") + group_size = required_int("conv_group_size") + rank = required_int("selector_rank") + top_k = required_int("selector_top_k") + + std = getattr(config, "initializer_range", 0.02) + + # Replace each layer's no-op sublayer wrappers with real convolutions. The + # backbone layer forward already calls prepare()/finish() around attention + # and the MLP, so nothing else in the layer changes. + for layer in self.layers: + for wrapper_name in ("attention_conv", "mlp_conv"): + setattr( + layer, + wrapper_name, + DFlashGroupedConv( + hidden_size=config.hidden_size, + block_size=self.block_size, + taps=taps, + group_size=group_size, + ), + ) + + self.candidate_selector = CandidateSelector( + hidden_size=config.hidden_size, + vocab_size=config.vocab_size, + rank=rank, + top_k=top_k, + std=std, + ) + + # DFlashModule.__init__ already ran _init_weights before these modules + # existed, so initialize the new Linear layers explicitly. base_kernel and + # the codebooks keep the init set in their own constructors. + self._init_head_weights(std) + + def _init_head_weights(self, std: float): + """Initialize the convolution and selector Linear layers (matching HF _init_weights).""" + linears = [self.candidate_selector.hidden_projection] + for layer in self.layers: + linears += [layer.attention_conv.kernel_projection, layer.mlp_conv.kernel_projection] + for module in linears: + nn.init.normal_(module.weight, mean=0.0, std=std) + if module.bias is not None: + nn.init.zeros_(module.bias) From 9975a5bc7db4c00164917f518a4947b39f6cd67a Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 19 Aug 2026 08:55:44 +0000 Subject: [PATCH 03/10] test(speculative): CPU unit tests for the DFlash2 draft variant Cover conversion routing, the convolution's structural invariants, the selector objective and the export contract. Two of these are the ones worth keeping: the convolution must be an identity at initialization (so enabling DFlash2 extends a DFlash backbone rather than perturbing it) and its taps must stay inside the block while still letting a position see its predecessors. A single-batch overfit guards the selector's target/predecessor alignment, which a finite decreasing loss alone does not. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../speculative/plugins/test_hf_dflash2.py | 362 ++++++++++++++++++ 1 file changed, 362 insertions(+) create mode 100644 tests/unit/torch/speculative/plugins/test_hf_dflash2.py diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py new file mode 100644 index 00000000000..b2ad7528665 --- /dev/null +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py @@ -0,0 +1,362 @@ +# 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. + +"""CPU unit tests for the DFlash2 speculative decoding plugin. + +DFlash2 reuses the DFlash mode/pipeline and adds grouped dynamic convolutions +around every attention/MLP sublayer plus a low-rank candidate selector. These +tests cover conversion routing, the convolution's two structural invariants +(identity at initialization, no leakage across the block boundary), the selector +training objective, and the export format against the SGLang/vLLM +``DFlash2DraftModel`` layout (``attention_conv.*`` / ``mlp_conv.*`` / +``candidate_selector.*``). +""" + +import json +from copy import deepcopy + +import pytest +import torch +from _test_utils.torch.transformers_models import get_tiny_llama +from safetensors.torch import load_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_dflash2 import HFDFlash2Model +from modelopt.torch.speculative.plugins.modeling_dflash import ( + DFlashModule, + _IdentitySublayerWrapper, +) +from modelopt.torch.speculative.plugins.modeling_dflash2 import DFlash2Module, DFlashGroupedConv + +BLOCK_SIZE = 4 +NUM_DRAFT_LAYERS = 2 +SEQ_LEN = 16 # must be a multiple of BLOCK_SIZE +CONV_KERNEL_SIZE = 2 +CONV_GROUP_SIZE = 4 +SELECTOR_RANK = 8 +SELECTOR_TOP_K = 5 + +ARCH_FIELDS = ["conv_kernel_size", "conv_group_size", "selector_rank", "selector_top_k"] + + +def _get_dflash2_config(selector_loss_alpha=1.0, block_size=BLOCK_SIZE, **arch_overrides): + """Create a DFlash2 config for testing (dflash mode + projector_type=dflash2).""" + config = deepcopy(DFLASH_DEFAULT_CFG["config"]) + config["dflash_block_size"] = block_size + config["dflash_use_torch_compile"] = False + config["dflash_mask_token_id"] = 0 # token 0 as mask for the tiny model + config["dflash_self_logit_distillation"] = False + config["dflash_selector_loss_alpha"] = selector_loss_alpha + config["dflash_architecture_config"] = { + "num_hidden_layers": NUM_DRAFT_LAYERS, + "projector_type": "dflash2", + "conv_kernel_size": CONV_KERNEL_SIZE, + "conv_group_size": CONV_GROUP_SIZE, + "selector_rank": SELECTOR_RANK, + "selector_top_k": SELECTOR_TOP_K, + **arch_overrides, + } + return config + + +def _make_batch(vocab_size): + torch.manual_seed(0) + input_ids = torch.randint(1, vocab_size, (2, SEQ_LEN)) + return input_ids, torch.ones_like(input_ids), input_ids.clone() + + +class TestDFlash2Convert: + """Test DFlash2 conversion routing and module construction.""" + + def test_convert_creates_dflash2_model(self): + """projector_type=dflash2 routes to HFDFlash2Model (a HFDFlashModel subclass).""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + assert isinstance(model, HFDFlash2Model) + assert isinstance(model, HFDFlashModel) + assert isinstance(model.dflash_module, DFlash2Module) + + def test_every_sublayer_wrapped_in_a_convolution(self): + """Both sublayer wrappers on every draft layer become real convolutions.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + layers = model.dflash_module.layers + assert len(layers) == NUM_DRAFT_LAYERS + for layer in layers: + for conv in (layer.attention_conv, layer.mlp_conv): + assert isinstance(conv, DFlashGroupedConv) + assert conv.taps == CONV_KERNEL_SIZE + assert conv.group_size == CONV_GROUP_SIZE + + def test_selector_shapes(self): + """The candidate selector's codebooks and projection are sized from the config.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + selector = model.dflash_module.candidate_selector + vocab = model.dflash_config.vocab_size + assert selector.top_k == SELECTOR_TOP_K + assert selector.predecessor_codebook.shape == (vocab, SELECTOR_RANK) + assert selector.successor_codebook.shape == (vocab, SELECTOR_RANK) + assert selector.hidden_projection.out_features == SELECTOR_RANK + assert selector.hidden_projection.bias is None + + def test_new_params_trainable(self): + """The convolution and selector parameters are trainable.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + new = [ + (n, p) + for n, p in model.named_parameters() + if "_conv." in n or "candidate_selector" in n + ] + assert len(new) >= 2 * NUM_DRAFT_LAYERS * 2 + 3 + assert all(p.requires_grad for _, p in new) + + @pytest.mark.parametrize("field", ARCH_FIELDS) + def test_missing_architecture_field_raises(self, field): + """projector_type=dflash2 without a required architecture field is an error.""" + config = _get_dflash2_config() + del config["dflash_architecture_config"][field] + model = get_tiny_llama(num_hidden_layers=4) + with pytest.raises(ValueError, match=field): + mtsp.convert(model, [("dflash", config)]) + + def test_conv_kernel_larger_than_block_raises(self): + """A convolution tap count exceeding the block size is an error.""" + model = get_tiny_llama(num_hidden_layers=4) + config = _get_dflash2_config(conv_kernel_size=BLOCK_SIZE + 1) + with pytest.raises(ValueError, match="conv_kernel_size"): + mtsp.convert(model, [("dflash", config)]) + + def test_conv_group_size_must_divide_hidden(self): + """A conv_group_size that does not divide hidden_size is an error.""" + model = get_tiny_llama(num_hidden_layers=4) + config = _get_dflash2_config(conv_group_size=model.config.hidden_size - 1) + with pytest.raises(ValueError, match="conv_group_size"): + mtsp.convert(model, [("dflash", config)]) + + def test_dflash_mode_still_creates_plain_dflash(self): + """Without projector_type=dflash2, conversion still yields a plain DFlash model.""" + config = deepcopy(DFLASH_DEFAULT_CFG["config"]) + config["dflash_mask_token_id"] = 0 + config["dflash_architecture_config"] = {"num_hidden_layers": NUM_DRAFT_LAYERS} + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config)]) + assert isinstance(model, HFDFlashModel) + assert not isinstance(model, HFDFlash2Model) + assert type(model.dflash_module) is DFlashModule + # The sublayer seam stays a parameterless no-op for a plain DFlash draft. + for layer in model.dflash_module.layers: + assert isinstance(layer.attention_conv, _IdentitySublayerWrapper) + assert isinstance(layer.mlp_conv, _IdentitySublayerWrapper) + assert not any("_conv." in n for n, _ in model.dflash_module.named_parameters()) + + +class TestDFlashGroupedConv: + """Test the convolution's structural invariants directly.""" + + def _conv(self, hidden_size=32, taps=2): + torch.manual_seed(0) + return DFlashGroupedConv( + hidden_size=hidden_size, block_size=BLOCK_SIZE, taps=taps, group_size=CONV_GROUP_SIZE + ).double() + + def test_identity_at_initialization(self): + """With the dynamic kernel zeroed, the identity base kernel is a no-op. + + This is what makes enabling DFlash2 a stable extension of a DFlash backbone + rather than a perturbation of it. + """ + conv = self._conv() + with torch.no_grad(): + conv.kernel_projection.weight.zero_() + x = torch.randn(2, SEQ_LEN, 32, dtype=torch.double) + out = conv.finish(*conv.prepare(x)) + assert torch.allclose(out, x, atol=1e-12) + + def test_taps_do_not_cross_the_block_boundary(self): + """Perturbing the last position of a block leaves later blocks untouched.""" + conv = self._conv() + x = torch.randn(2, SEQ_LEN, 32, dtype=torch.double) + baseline = conv.finish(*conv.prepare(x)) + + perturbed_input = x.clone() + perturbed_input[:, BLOCK_SIZE - 1] += 5.0 + perturbed = conv.finish(*conv.prepare(perturbed_input)) + + assert torch.allclose(baseline[:, BLOCK_SIZE:], perturbed[:, BLOCK_SIZE:], atol=1e-12) + assert not torch.allclose(baseline[:, :BLOCK_SIZE], perturbed[:, :BLOCK_SIZE]) + + def test_intra_block_dependency_is_backward_only(self): + """A position influences its successors inside the block, never its predecessors. + + This is the point of the convolution: it injects the sequential dependency the + parallel backbone lacks, without letting a position see the future. + """ + conv = self._conv() + x = torch.randn(2, SEQ_LEN, 32, dtype=torch.double) + baseline = conv.finish(*conv.prepare(x)) + + perturbed_input = x.clone() + perturbed_input[:, 1] += 5.0 + perturbed = conv.finish(*conv.prepare(perturbed_input)) + + assert torch.allclose(baseline[:, 0], perturbed[:, 0], atol=1e-12) + assert not torch.allclose(baseline[:, 2], perturbed[:, 2]) + + def test_sequence_length_must_be_block_aligned(self): + """A sequence length not divisible by the block size is an error.""" + conv = self._conv() + with pytest.raises(ValueError, match="block_size"): + conv.prepare(torch.randn(1, BLOCK_SIZE + 1, 32, dtype=torch.double)) + + +class TestDFlash2Forward: + """Test the DFlash2 training forward (online path on CPU).""" + + def test_forward_grads_reach_conv_and_selector(self): + """Backward fills gradients on the convolutions, the selector and the backbone.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + model.train() + + input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) + out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + assert out.loss.requires_grad + assert out.loss.dim() == 0 + out.loss.backward() + + module = model.dflash_module + for grad in ( + module.candidate_selector.predecessor_codebook.grad, + module.candidate_selector.successor_codebook.grad, + module.layers[0].attention_conv.base_kernel.grad, + module.layers[0].mlp_conv.kernel_projection.weight.grad, + module.fc.weight.grad, + ): + assert grad is not None and torch.isfinite(grad).all() + assert grad.abs().sum() > 0 + + def test_selector_metrics_reported(self): + """The forward records selector accuracy and top-k coverage.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + model.train() + + input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) + model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + metrics = model._selector_metrics + for key in ("selector_accuracy", "selector_coverage"): + assert 0.0 <= metrics[key] <= 1.0 + + def test_selector_alpha_zero_disables_the_term(self): + """alpha=0 trains the backbone and convolutions only; the selector gets no grad.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config(selector_loss_alpha=0.0))]) + model.train() + + input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) + out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + out.loss.backward() + + module = model.dflash_module + codebook_grad = module.candidate_selector.predecessor_codebook.grad + assert codebook_grad is None or codebook_grad.abs().sum() == 0 + # The convolutions still train: they live inside the backbone. + conv_grad = module.layers[0].attention_conv.kernel_projection.weight.grad + assert conv_grad is not None and conv_grad.abs().sum() > 0 + + def test_selector_loss_increases_total_loss(self): + """The selector term adds to the backbone loss rather than replacing it.""" + input_ids, attention_mask, labels = _make_batch(32) + + losses = {} + for alpha in (0.0, 1.0): + torch.manual_seed(0) + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config(selector_loss_alpha=alpha))]) + model.train() + torch.manual_seed(0) + out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + losses[alpha] = float(out.loss.detach()) + assert losses[1.0] > losses[0.0] + + def test_overfits_a_single_batch(self): + """A few steps on one batch drive backbone and selector accuracy up. + + Guards the target/predecessor alignment: a misaligned selector objective still + produces a finite decreasing loss, but its accuracy does not reach 1. + """ + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + model.train() + + input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) + optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=5e-3) + for _ in range(60): + out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + optimizer.zero_grad() + out.loss.backward() + optimizer.step() + + assert out.train_acc[0][0] > 0.9 + assert model._selector_metrics["selector_accuracy"] > 0.9 + + +class TestDFlash2Export: + """Test the DFlash2 export format (weights + config).""" + + def _export(self, tmp_path): + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + export_dir = tmp_path / "exported" + model.get_exporter().export(export_dir) + return export_dir + + def test_export_weight_keys_match_reference(self, tmp_path): + """Exported weights carry the DFlash2 tensors under reference names, no prefix.""" + sd = load_file(str(self._export(tmp_path) / "model.safetensors")) + for key in sd: + assert "dflash_module." not in key + assert "rotary_emb" not in key + + assert "candidate_selector.predecessor_codebook" in sd + assert "candidate_selector.successor_codebook" in sd + assert "candidate_selector.hidden_projection.weight" in sd + for layer_idx in range(NUM_DRAFT_LAYERS): + for wrapper in ("attention_conv", "mlp_conv"): + assert f"layers.{layer_idx}.{wrapper}.base_kernel" in sd + assert f"layers.{layer_idx}.{wrapper}.kernel_projection.weight" in sd + + def test_export_config_declares_dflash2_architecture(self, tmp_path): + """config.json selects the DFlash2 serving path and carries its fields. + + The architecture name matters: a checkpoint declaring ``DFlashDraftModel`` + loads as a plain DFlash draft and silently ignores these weights. + """ + with open(self._export(tmp_path) / "config.json") as f: + cfg = json.load(f) + + assert cfg["architectures"] == ["DFlash2DraftModel"] + dflash_config = cfg["dflash_config"] + assert dflash_config["projector_type"] == "dflash2" + assert dflash_config["conv_kernel_size"] == CONV_KERNEL_SIZE + assert dflash_config["conv_group_size"] == CONV_GROUP_SIZE + assert dflash_config["selector_rank"] == SELECTOR_RANK + assert dflash_config["selector_top_k"] == SELECTOR_TOP_K + assert "mask_token_id" in dflash_config + assert "target_layer_ids" in dflash_config From 197dd851cfd4df8e58f7f9ea409ed07441eb66e7 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 19 Aug 2026 08:57:17 +0000 Subject: [PATCH 04/10] docs(changelog): note the DFlash2 draft variant Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- CHANGELOG.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4580a045f39..6253d17f37e 100755 --- a/CHANGELOG.rst +++ b/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. From 13e8e2f8e60891017fa9665b3b137d70845fec4d Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 19 Aug 2026 09:34:57 +0000 Subject: [PATCH 05/10] feat(recipes): DFlash2 training recipe Mirrors dspark.yaml with the DFlash2 architecture fields (conv taps/group size, selector rank/top-k) and dflash_selector_loss_alpha in place of the DSpark head and its three-term loss weights. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../general/speculative_decoding/dflash2.yaml | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 modelopt_recipes/general/speculative_decoding/dflash2.yaml diff --git a/modelopt_recipes/general/speculative_decoding/dflash2.yaml b/modelopt_recipes/general/speculative_decoding/dflash2.yaml new file mode 100644 index 00000000000..86555f357f9 --- /dev/null +++ b/modelopt_recipes/general/speculative_decoding/dflash2.yaml @@ -0,0 +1,98 @@ +# DFlash2 speculative-decoding training recipe. +# +# DFlash2 (https://inco.ai/blog/dflash2/) reuses the DFlash mode/pipeline and adds +# two components, selected via dflash_architecture_config.projector_type=dflash2: +# - a grouped dynamic depthwise convolution around every attention/MLP sublayer, +# giving each block position a view of its predecessors inside the block; +# - a low-rank candidate selector that scores transitions between adjacent block +# positions' top-k candidates, so serving walks one coherent path. +# The selector is trained by an extra cross-entropy term weighted by +# dflash_selector_loss_alpha. Online training is the default path (data.mode=online). +# Override fields via an OmegaConf dotlist. + +metadata: + recipe_type: speculative_dflash + description: DFlash2 training recipe (DFlash backbone + sublayer conv + candidate selector). + +# maps to ModelArguments (main.py) +model: + model_name_or_path: + trust_remote_code: false + use_fake_base_for_offline: false + +# maps to DataArguments (main.py) +data: + mode: online + data_path: + offline_data_path: + # Jinja chat template with {% generation %} tags for answer_only_loss. + chat_template: + +# maps to TrainingArguments (main.py) +training: + # --- commonly modified --- + output_dir: + num_train_epochs: 6 + per_device_train_batch_size: 1 + learning_rate: 6.0e-4 + warmup_ratio: 0.04 + training_seq_len: 3072 + logging_steps: 50 + save_steps: 2000 + cp_size: 1 + dp_shard_size: 1 + disable_tqdm: true + # Keep off: eval runs the DFlash backbone only (Markov head not applied yet), + # so AR would reflect the backbone alone, not the trained model. Compare via + # export + the offline acceptance-length harness instead. + estimate_ar: false + ar_validate_steps: 0 + answer_only_loss: true + + # --- rarely modified --- + 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 + # Safe default: the selector params are unused when + # dflash_selector_loss_alpha == 0, which would otherwise trip DDP. + ddp_find_unused_parameters: true + ddp_timeout: 1800 + report_to: tensorboard + +# maps to DFlashConfig (modelopt/torch/speculative/config.py). +dflash: + dflash_block_size: 16 + dflash_num_anchors: 256 + dflash_use_torch_compile: false + dflash_self_logit_distillation: false + # gamma for exponential loss decay (block_size=16 -> 7). + dflash_loss_decay_factor: 7.0 + # Qwen3 has no native mask token; 151669 is an unused id used by the reference. + dflash_mask_token_id: 151669 + # Weight of the candidate-selector cross-entropy term (0 disables it and trains + # the backbone + convolutions only). + dflash_selector_loss_alpha: 1.0 + dflash_architecture_config: + num_hidden_layers: 5 + # Draft attention/MLP dims — set explicitly (the draft is an independent + # Qwen3 model and does NOT inherit these from the base). GQA: 8 KV heads. + num_attention_heads: 32 + num_key_value_heads: 8 + head_dim: 128 + intermediate_size: 12288 + projector_type: dflash2 + # Grouped dynamic depthwise convolution. conv_kernel_size is the number of taps + # (2 = each position also sees its predecessor); it must not exceed the block + # size. conv_group_size must divide hidden_size. + conv_kernel_size: 2 + conv_group_size: 16 + # Candidate selector: rank of the transition codebooks, and how many of the + # backbone top-k candidates per position it re-ranks. + selector_rank: 256 + selector_top_k: 16 From ba377e7a355c7e7de83c6472d09c3475706407af Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 19 Aug 2026 12:42:20 +0000 Subject: [PATCH 06/10] fix(speculative): read rope_theta from rope_parameters first A Transformers 5 config can carry BOTH a top-level rope_theta and a rope_parameters dict holding a different value: the real base lives in rope_parameters while the config class default (10000.0 for Qwen3) stays visible as the flat attribute. Reading the flat field first therefore picked up 10000.0 for a Qwen3-8B target whose actual base is 1000000. DFlash injects the target's KV into every draft layer, so a draft built this way trains, exports and loads without complaint while its RoPE base is 100x off the target's. Observed on an NRT smoke: the exported draft carried rope_theta 10000.0 where the reference z-lab checkpoint has 1000000, and vLLM died during engine warmup. Prefer rope_parameters in both the exporter's _get_rope_theta and the training-side enforcement in HFDFlashModel.modify, and keep the draft's own rope_parameters dict in sync with the flat field it is derived from. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 139 ++++++++++++++++++ .../torch/export/plugins/hf_spec_export.py | 18 ++- .../torch/speculative/plugins/hf_dflash.py | 21 ++- 3 files changed, 171 insertions(+), 7 deletions(-) create mode 100644 modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md new file mode 100644 index 00000000000..fddc925ce5d --- /dev/null +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -0,0 +1,139 @@ +# Export: Model-Specific Logic Refactor — Action Plan + +**Goal:** Give the unified HF export path a per-model data registry +(top-level `modelopt/models/`, one file per HF model type), so that supporting a +new model means adding one declarative spec file instead of editing if/elif chains +across the export engine. + +**Scope:** The unified HF export path (`unified_export_hf.py` + its helpers) only. + +- The **Megatron** path (`plugins/mcore_*`) already follows a per-model registry + pattern and is untouched. +- The **TRT-LLM** checkpoint path (`model_config_export.py`, the `build_*` + functions in `layer_utils.py`, `tensorrt_llm_utils.py`) is legacy: the plan is + to move it out as-is (separate track), not to refactor it. Its per-model + branches stay where they are. + +## 1. Where the HF path stands after PR #1939 + +PR #1939 gave the HF path registry-based **module dispatch**: `ExportModuleRegistry` +and `PrepareMoEInputsRegistry` (`registry.py`, `hf_export_handlers.py`) select +*which handler processes a module* by class/predicate match, replacing the if/elif +chains that used to live in `unified_export_hf.py`. + +What is still hardcoded is the **per-model data** those handlers (and other HF-path +helpers) consume. Inventory of model-specific logic reachable from the HF path: + +| Item | Location | Kind | +|---|---|---| +| MoE expert linear names (Qwen/DeepSeek/Mixtral/DBRX/GptOss/NemotronH/Gemma4) | `layer_utils.get_expert_linear_names` | data — **migrated (P1)** | +| Duplicate expert-naming table + iterable-experts support gate | `layer_utils.get_experts_list` | data — **migrated (P1)** | +| MoE block class-name list | `layer_utils.is_moe` | data — **migrated (P3)** | +| AWQ `pre_quant_scale` fusion rules (Llama/Qwen3) | `quant_utils.PQS_FUSE_MODULE_MAPPING` | data — **migrated (P2)** | +| weight+1 layernorm class names (Gemma RMSNorm, LayerNorm1P) | `quant_utils._layernorm_uses_weight_plus_one` | data — **migrated (NormSpec)** | +| MoE gate/up fusion pairs (`_GATE_UP_PAIRS`, also privately imported by `quantization/model_calib.py`) | `layer_utils.sync_moe_gate_up_amax` | data — **migrated (MoESpec.gate_up_pair)** | +| BMM-style expert class list (`Llama4TextExperts`/`GptOssExperts`) inline copy for weight transpose | `quant_utils` (dispatch copies in `hf_export_handlers.py` stay) | data — P4 | +| VLM detection gates (`phi4mm` model_type, `nemotronparse` architecture, `"nemotron" in model_type` tower special case) | `model_utils.py`, `unified_export_hf.py` | data (gates) + behavior (extraction) — collect until worth a spec flag | +| Handler match keys (Llama4TextExperts, GptOssExperts, DbrxExperts, QuantMoELinear) | `hf_export_handlers.py` | dispatch (stays: structural, per-module) | +| Fused-expert gated/non-gated split (`gate_up_proj` vs `up_proj`) | `moe_utils.py` | data + structure | +| dummy-forward special cases (Whisper input, Nemotron-VL tower) | `unified_export_hf.requantize_resmooth_fused_llm_layers` | behavior | +| VLM language-tower extraction, DiffusionGemma tied-key reorder | `model_utils.py` | behavior | + +## 2. Design + +Two layers: + +- **Engine** (existing files, organized by operation): owns all algorithms and the + module walk; consults the registry for per-model values. +- **Modeling library** (top-level `modelopt/models/`, one file per HF model type): + declarative per-model **data only**. No export logic, stdlib-only imports (not even + torch), so it sits at the bottom of the dependency graph and any modelopt subsystem + can depend on it. + +```text +modelopt/models/ + specs.py # ModelSpec — the ONE global per-model descriptor, composed + # from section mixins: topic sections (MoESpec: MoE layouts as + # MoEVariant tuples; NormSpec) + subsystem sections (ExportSpec); + # future sections mix in the same way + registry.py # register() + lookups queried by spec type (None when unmatched) + # + the MRO exact-name matching core (match_class_names) + __init__.py # re-exports; importing it registers all specs + .py # one small file per HF model type (mirrors + # transformers.models); import == registration +``` + +Model type names mirror +[`transformers.models`](https://github.com/huggingface/transformers/tree/main/src/transformers/models) +(e.g. `qwen3_moe.py`, `gpt_oss.py`, `nemotron_h.py`); trust-remote-code models +(`arctic`, `deepseek`) use their config `model_type`. + +Each model registers exactly ONE `ModelSpec` (registry enforces uniqueness), so +`get_spec(model_type)` is a dict lookup and consumers call methods directly on the +global spec (e.g. `spec.expert_linear_names_for(module)`). +Resolution is by model type, mirroring HF's own indexing: the engine reads the +model's root HF type (`model.config.model_type`) once per export and passes it +down. The lookup is strict: only the model's own spec is consulted, so a model +whose model_type has no spec fails loudly (register a spec) instead of inheriting +a neighbor's data through a coincidental class-name match. Sub-config model types +of composite models are not walked today — a composite whose MoE lives under a +tower type registers the root type too (`gemma4` + `gemma4_text`, following the +`gemma3`/`gemma3_text` precedent); a recursive config walk is a future extension +if a real VLM tower needs it. Within the spec, each `MoESpec` nests one +`MoEVariant` per concrete block layout — several when the same checkpoint +materializes with different classes and projection names (Mixtral across +transformers generations); variant `block_names` (matched against the module +MRO) identify MoE blocks (`is_moe`) and pick the variant. +`get_expert_linear_names` doesn't need the block class at all when a model's +variants agree on one naming. Without a model type (no config available: unit +tests, the TRT-LLM path), lookups search all specs by class name. + +During migration, call sites kept the legacy branches as a fallback behind the +spec lookup. Once the specs covered every family the legacy chains served, the +chains — and the silent ``w1/w2/w3`` guess for unknown models — were deleted: +expert-name resolution is now *structural detection -> spec -> raise*, so a new +MoE model fails loudly, asking for a spec, instead of inheriting another model's +naming. Generic detection that is not per-model data (the ``*SparseMoeBlock`` / +``*MoeLayer`` conventions and the router+experts structural check in ``is_moe``, +the fused-experts quantizer probe in ``get_expert_linear_names``) stays in the +engine, ahead of or beside the spec lookup. + +Note on naming: `modelopt/models/registry.py` (per-model **data**, "what are +this model's values") is distinct from the export-path `registry.py` from PR #1939 +(per-module **dispatch**, "which handler processes this module"). The two layers +compose: handlers look up model data through `modelopt.models`. + +## 3. Migration plan + +Each step is one PR with a fallback to the legacy path and an equivalence check +against existing export tests. + +| Step | What | Status | +|---|---|---| +| **P1** | Registry skeleton + MoE expert naming: `get_expert_linear_names` and `get_experts_list` read `spec.expert_linear_names` / `spec.has_iterable_experts`. The #1 "add a MoE model" shotgun-surgery driver. | this PR | +| **P2** | `PQS_FUSE_MODULE_MAPPING` → `spec.pqs_fuse_rules`, aggregated via `iter_pqs_fuse_rules` (llama/qwen3 specs). | this PR | +| **P3** | `is_moe` explicit class-name list → `spec.moe_block_names` (arctic/dbrx_ffn are identification-only specs: no expert naming, so expert-name lookups keep the engine default). The generic `*SparseMoeBlock`/`*MoeLayer` conventions and the structural router+experts check stay in the engine. | this PR | +| **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs; share the matcher machinery with the export dispatch registry (#1939). Model_type-scoped resolution (`collect_model_types` + scoped `match_moe_block`, threaded via `ExportContext.model_types`) is already in place from this PR. | planned | +| **P5** | Cross-subsystem pilot: unify the remaining copies of linear-fusion-group knowledge (see §4) into spec fields. Partially done: `_GATE_UP_PAIRS` became `MoESpec.gate_up_pair` and `model_calib.py`'s private import of it is gone — the first quantization consumer of `modelopt.models`. Remaining: `shared_input.SHARED_PATTERNS`, `algorithms.quant_grouping_rules`. | in progress | +| **P6** | Migrate remaining quantization-side data: default disabled-quantizer patterns, on-the-fly conversion gates, AutoQuantize grouping rules (see §4). | planned | +| **OUT** | TRT-LLM path branches (`decoder_type` chains in `build_*`, `model_config_export.py`, `tensorrt_llm_utils.py`): frozen, moved out unchanged on a separate track. Candidates for deletion on that track: `adjust_attn_amax_values`, `update_experts_avg_prequant_scale` (unused). NOTE: `MODEL_NAME_TO_TYPE` / `get_model_type` are NOT dead — `examples/hf_ptq/hf_ptq.py` and `multinode_ptq.py` still call them; migrate the examples before removing. | separate track | + +**Guardrails:** one data category per PR (fallback-first while a category is +partially migrated, explicit-error once specs cover it); the engine keeps the +algorithms — model specs supply values only, never fork functions. + +## 4. Beyond export: per-model data in quantization (P5/P6 inventory) + +The same three kinds of model-specific logic exist on the quantization side. Only +kind (a) migrates into `modelopt/modeling`; (b) stays in each subsystem's module +registry (a spec may hold pointer data, never the surgery code); (c) stays in the +engine behind structural checks. + +| Item | Location | Kind | +|---|---|---| +| Linear fusion groups (q/k/v, gate/up, `w1/w3`) — was duplicated 3x; the `_GATE_UP_PAIRS` copy is now `MoESpec.gate_up_pair` | `quantization/utils/shared_input.SHARED_PATTERNS`, `quantization/algorithms.quant_grouping_rules` (remaining) | data — P5 | +| Model-class gates for on-the-fly conversion (`"DbrxForCausalLM"`, `("Step3p5ForCausalLM", ...)`) | `quantization/plugins/huggingface.py` | data — P6 | +| Default disabled-quantizer patterns (`*router*`, `*vision_tower*`; per-model, NVBug-gated) | `modelopt_recipes/.../default_disabled_quantizers.yaml` | data — P6 | +| AutoQuantize grouping regexes (llama q/k/v, Mixtral `w1/w2/w3`, NemotronH mixer) | `quantization/algorithms.py` | data — P6 | +| Quant wrapper classes (`_QuantDbrxExperts` splits `w1/v1/w2` into per-expert linears) | `quantization/plugins/huggingface.py` via `QuantModuleRegistry` | dispatch (stays) | +| Structural MoE detection (`gate`+`experts`+`top_k` attrs; 3-D `gate_up_proj` -> gated) | `quantization/plugins/huggingface.py` | behavior (stays) | diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 820c5d58f53..01f771cb454 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -31,11 +31,15 @@ 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"): @@ -43,6 +47,10 @@ def _get_rope_theta(config, default=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 diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index fe13f9c6b2a..0b146a39064 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/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, From fe8d5b8d848cd20cae9e1e28d8f41debe2605a8e Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:21:14 +0000 Subject: [PATCH 07/10] feat(export): match the published DFlash2 config contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The released DFlash2 checkpoints (z-lab/Qwen3.8-27B-DFlash2) carry block_size inside dflash_config and state is_causal explicitly rather than leaving it to be inferred from layer_types. Emit both. is_causal matters because vLLM's _dflash_layer_causal falls back to `layer_types[i] == "sliding_attention"`, which reads a sliding-window draft as causal — the published checkpoints override that with an explicit false. A full-attention draft resolved to the same value already, so this pins existing behaviour rather than changing it. Verified against the released checkpoint on NRT: 81 tensors, 21 name patterns, zero difference in either direction. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/torch/export/plugins/hf_spec_export.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 01f771cb454..5b8b81016db 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -572,6 +572,14 @@ def _export_config(self): "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 From c446f6e43607bf27a3f7e7609496382610bb3f8b Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:22:34 +0000 Subject: [PATCH 08/10] chore: drop MODEL_SPECIFIC_REFACTOR.md from this branch The file was picked up by a `git add -A` in an earlier commit on this branch; it belongs to unrelated in-flight work and is not part of the DFlash2 change. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 139 ------------------ 1 file changed, 139 deletions(-) delete mode 100644 modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md deleted file mode 100644 index fddc925ce5d..00000000000 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ /dev/null @@ -1,139 +0,0 @@ -# Export: Model-Specific Logic Refactor — Action Plan - -**Goal:** Give the unified HF export path a per-model data registry -(top-level `modelopt/models/`, one file per HF model type), so that supporting a -new model means adding one declarative spec file instead of editing if/elif chains -across the export engine. - -**Scope:** The unified HF export path (`unified_export_hf.py` + its helpers) only. - -- The **Megatron** path (`plugins/mcore_*`) already follows a per-model registry - pattern and is untouched. -- The **TRT-LLM** checkpoint path (`model_config_export.py`, the `build_*` - functions in `layer_utils.py`, `tensorrt_llm_utils.py`) is legacy: the plan is - to move it out as-is (separate track), not to refactor it. Its per-model - branches stay where they are. - -## 1. Where the HF path stands after PR #1939 - -PR #1939 gave the HF path registry-based **module dispatch**: `ExportModuleRegistry` -and `PrepareMoEInputsRegistry` (`registry.py`, `hf_export_handlers.py`) select -*which handler processes a module* by class/predicate match, replacing the if/elif -chains that used to live in `unified_export_hf.py`. - -What is still hardcoded is the **per-model data** those handlers (and other HF-path -helpers) consume. Inventory of model-specific logic reachable from the HF path: - -| Item | Location | Kind | -|---|---|---| -| MoE expert linear names (Qwen/DeepSeek/Mixtral/DBRX/GptOss/NemotronH/Gemma4) | `layer_utils.get_expert_linear_names` | data — **migrated (P1)** | -| Duplicate expert-naming table + iterable-experts support gate | `layer_utils.get_experts_list` | data — **migrated (P1)** | -| MoE block class-name list | `layer_utils.is_moe` | data — **migrated (P3)** | -| AWQ `pre_quant_scale` fusion rules (Llama/Qwen3) | `quant_utils.PQS_FUSE_MODULE_MAPPING` | data — **migrated (P2)** | -| weight+1 layernorm class names (Gemma RMSNorm, LayerNorm1P) | `quant_utils._layernorm_uses_weight_plus_one` | data — **migrated (NormSpec)** | -| MoE gate/up fusion pairs (`_GATE_UP_PAIRS`, also privately imported by `quantization/model_calib.py`) | `layer_utils.sync_moe_gate_up_amax` | data — **migrated (MoESpec.gate_up_pair)** | -| BMM-style expert class list (`Llama4TextExperts`/`GptOssExperts`) inline copy for weight transpose | `quant_utils` (dispatch copies in `hf_export_handlers.py` stay) | data — P4 | -| VLM detection gates (`phi4mm` model_type, `nemotronparse` architecture, `"nemotron" in model_type` tower special case) | `model_utils.py`, `unified_export_hf.py` | data (gates) + behavior (extraction) — collect until worth a spec flag | -| Handler match keys (Llama4TextExperts, GptOssExperts, DbrxExperts, QuantMoELinear) | `hf_export_handlers.py` | dispatch (stays: structural, per-module) | -| Fused-expert gated/non-gated split (`gate_up_proj` vs `up_proj`) | `moe_utils.py` | data + structure | -| dummy-forward special cases (Whisper input, Nemotron-VL tower) | `unified_export_hf.requantize_resmooth_fused_llm_layers` | behavior | -| VLM language-tower extraction, DiffusionGemma tied-key reorder | `model_utils.py` | behavior | - -## 2. Design - -Two layers: - -- **Engine** (existing files, organized by operation): owns all algorithms and the - module walk; consults the registry for per-model values. -- **Modeling library** (top-level `modelopt/models/`, one file per HF model type): - declarative per-model **data only**. No export logic, stdlib-only imports (not even - torch), so it sits at the bottom of the dependency graph and any modelopt subsystem - can depend on it. - -```text -modelopt/models/ - specs.py # ModelSpec — the ONE global per-model descriptor, composed - # from section mixins: topic sections (MoESpec: MoE layouts as - # MoEVariant tuples; NormSpec) + subsystem sections (ExportSpec); - # future sections mix in the same way - registry.py # register() + lookups queried by spec type (None when unmatched) - # + the MRO exact-name matching core (match_class_names) - __init__.py # re-exports; importing it registers all specs - .py # one small file per HF model type (mirrors - # transformers.models); import == registration -``` - -Model type names mirror -[`transformers.models`](https://github.com/huggingface/transformers/tree/main/src/transformers/models) -(e.g. `qwen3_moe.py`, `gpt_oss.py`, `nemotron_h.py`); trust-remote-code models -(`arctic`, `deepseek`) use their config `model_type`. - -Each model registers exactly ONE `ModelSpec` (registry enforces uniqueness), so -`get_spec(model_type)` is a dict lookup and consumers call methods directly on the -global spec (e.g. `spec.expert_linear_names_for(module)`). -Resolution is by model type, mirroring HF's own indexing: the engine reads the -model's root HF type (`model.config.model_type`) once per export and passes it -down. The lookup is strict: only the model's own spec is consulted, so a model -whose model_type has no spec fails loudly (register a spec) instead of inheriting -a neighbor's data through a coincidental class-name match. Sub-config model types -of composite models are not walked today — a composite whose MoE lives under a -tower type registers the root type too (`gemma4` + `gemma4_text`, following the -`gemma3`/`gemma3_text` precedent); a recursive config walk is a future extension -if a real VLM tower needs it. Within the spec, each `MoESpec` nests one -`MoEVariant` per concrete block layout — several when the same checkpoint -materializes with different classes and projection names (Mixtral across -transformers generations); variant `block_names` (matched against the module -MRO) identify MoE blocks (`is_moe`) and pick the variant. -`get_expert_linear_names` doesn't need the block class at all when a model's -variants agree on one naming. Without a model type (no config available: unit -tests, the TRT-LLM path), lookups search all specs by class name. - -During migration, call sites kept the legacy branches as a fallback behind the -spec lookup. Once the specs covered every family the legacy chains served, the -chains — and the silent ``w1/w2/w3`` guess for unknown models — were deleted: -expert-name resolution is now *structural detection -> spec -> raise*, so a new -MoE model fails loudly, asking for a spec, instead of inheriting another model's -naming. Generic detection that is not per-model data (the ``*SparseMoeBlock`` / -``*MoeLayer`` conventions and the router+experts structural check in ``is_moe``, -the fused-experts quantizer probe in ``get_expert_linear_names``) stays in the -engine, ahead of or beside the spec lookup. - -Note on naming: `modelopt/models/registry.py` (per-model **data**, "what are -this model's values") is distinct from the export-path `registry.py` from PR #1939 -(per-module **dispatch**, "which handler processes this module"). The two layers -compose: handlers look up model data through `modelopt.models`. - -## 3. Migration plan - -Each step is one PR with a fallback to the legacy path and an equivalence check -against existing export tests. - -| Step | What | Status | -|---|---|---| -| **P1** | Registry skeleton + MoE expert naming: `get_expert_linear_names` and `get_experts_list` read `spec.expert_linear_names` / `spec.has_iterable_experts`. The #1 "add a MoE model" shotgun-surgery driver. | this PR | -| **P2** | `PQS_FUSE_MODULE_MAPPING` → `spec.pqs_fuse_rules`, aggregated via `iter_pqs_fuse_rules` (llama/qwen3 specs). | this PR | -| **P3** | `is_moe` explicit class-name list → `spec.moe_block_names` (arctic/dbrx_ffn are identification-only specs: no expert naming, so expert-name lookups keep the engine default). The generic `*SparseMoeBlock`/`*MoeLayer` conventions and the structural router+experts check stay in the engine. | this PR | -| **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs; share the matcher machinery with the export dispatch registry (#1939). Model_type-scoped resolution (`collect_model_types` + scoped `match_moe_block`, threaded via `ExportContext.model_types`) is already in place from this PR. | planned | -| **P5** | Cross-subsystem pilot: unify the remaining copies of linear-fusion-group knowledge (see §4) into spec fields. Partially done: `_GATE_UP_PAIRS` became `MoESpec.gate_up_pair` and `model_calib.py`'s private import of it is gone — the first quantization consumer of `modelopt.models`. Remaining: `shared_input.SHARED_PATTERNS`, `algorithms.quant_grouping_rules`. | in progress | -| **P6** | Migrate remaining quantization-side data: default disabled-quantizer patterns, on-the-fly conversion gates, AutoQuantize grouping rules (see §4). | planned | -| **OUT** | TRT-LLM path branches (`decoder_type` chains in `build_*`, `model_config_export.py`, `tensorrt_llm_utils.py`): frozen, moved out unchanged on a separate track. Candidates for deletion on that track: `adjust_attn_amax_values`, `update_experts_avg_prequant_scale` (unused). NOTE: `MODEL_NAME_TO_TYPE` / `get_model_type` are NOT dead — `examples/hf_ptq/hf_ptq.py` and `multinode_ptq.py` still call them; migrate the examples before removing. | separate track | - -**Guardrails:** one data category per PR (fallback-first while a category is -partially migrated, explicit-error once specs cover it); the engine keeps the -algorithms — model specs supply values only, never fork functions. - -## 4. Beyond export: per-model data in quantization (P5/P6 inventory) - -The same three kinds of model-specific logic exist on the quantization side. Only -kind (a) migrates into `modelopt/modeling`; (b) stays in each subsystem's module -registry (a spec may hold pointer data, never the surgery code); (c) stays in the -engine behind structural checks. - -| Item | Location | Kind | -|---|---|---| -| Linear fusion groups (q/k/v, gate/up, `w1/w3`) — was duplicated 3x; the `_GATE_UP_PAIRS` copy is now `MoESpec.gate_up_pair` | `quantization/utils/shared_input.SHARED_PATTERNS`, `quantization/algorithms.quant_grouping_rules` (remaining) | data — P5 | -| Model-class gates for on-the-fly conversion (`"DbrxForCausalLM"`, `("Step3p5ForCausalLM", ...)`) | `quantization/plugins/huggingface.py` | data — P6 | -| Default disabled-quantizer patterns (`*router*`, `*vision_tower*`; per-model, NVBug-gated) | `modelopt_recipes/.../default_disabled_quantizers.yaml` | data — P6 | -| AutoQuantize grouping regexes (llama q/k/v, Mixtral `w1/w2/w3`, NemotronH mixer) | `quantization/algorithms.py` | data — P6 | -| Quant wrapper classes (`_QuantDbrxExperts` splits `w1/v1/w2` into per-expert linears) | `quantization/plugins/huggingface.py` via `QuantModuleRegistry` | dispatch (stays) | -| Structural MoE detection (`gate`+`experts`+`top_k` attrs; 3-D `gate_up_proj` -> gated) | `quantization/plugins/huggingface.py` | behavior (stays) | From bf361100fe740db434acfd03d142e9f59db39087 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Sun, 23 Aug 2026 14:41:44 +0000 Subject: [PATCH 09/10] chore(pre-commit): exempt the DFlash2 plugins from the license-header hook MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit hf_dflash2.py and modeling_dflash2.py carry an upstream copyright header, so the hook must not prepend ours above it — same treatment the dflash / domino / dspark plugins already get in this list. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .pre-commit-config.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 85b5577494c..ea4814557da 100644 --- a/.pre-commit-config.yaml +++ b/.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| From fd085f6d919e1254850c9c1bb8249c4661e89999 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Sun, 23 Aug 2026 14:41:58 +0000 Subject: [PATCH 10/10] feat(launcher): Qwen3-8B DFlash2 online training example Mirrors the existing hf_online_dflash / hf_online_domino examples so the three correction heads are launched the same way, and documents the deltas needed to reproduce the published Qwen3-8B DFlash2 curve (3 epochs over Spec-Decoding-Dataset-v1 on 8x8 H100) rather than the short convergence check the file runs as configured. The five draft dims are spelled out explicitly: a dim left unset falls back to the Qwen3Config default instead of the base model's, which would silently train a different-sized drafter and void the comparison against the other arms. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../Qwen/Qwen3-8B/hf_online_dflash2.yaml | 107 ++++++++++++++++++ 1 file changed, 107 insertions(+) create mode 100644 tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml new file mode 100644 index 00000000000..e300e0665b4 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml @@ -0,0 +1,107 @@ +# DFlash2 online speculative decoding training for Qwen3-8B. +# +# DFlash2 = the DFlash draft backbone plus two additions: +# * a grouped dynamic depthwise convolution wrapped around every attention and +# MLP sublayer (taps clipped at block boundaries, identity-initialized so a +# fresh DFlash2 draft computes exactly what its DFlash backbone would), and +# * a low-rank candidate selector that scores transitions between adjacent +# block positions' top-k candidates, so serving walks one coherent path. +# See the dflash2.yaml recipe and +# modelopt/torch/speculative/plugins/{modeling,hf}_dflash2.py. +# +# 2-step pipeline: +# task_0: Build training conversations (Daring-Anteater multi-turn SFT, 50K) +# task_1: Online DFlash2 training + export of the drafter checkpoint +# +# As configured this is a short convergence check (max_steps=2000), matching the +# other Qwen3-8B online examples so it finishes on one node. To reproduce the +# published Qwen3-8B DFlash2 curve instead, see "Full run" below. +# +# Reference: inco.ai/blog/dflash2 | vLLM PR #52816 (serving support) +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml --yes +# uv run slurm.py --yaml modules/Model-Optimizer/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml --yes +# +# Full run (the strict A/B against DFlash / DSpark / Domino, 3 epochs = ~92K +# steps on the 1.96M-conversation Spec-Decoding-Dataset-v1, 8 nodes x 8 H100): +# - point data.data_path at that corpus instead of task_0's output +# - training.num_train_epochs=3 and drop training.max_steps +# - training.save_steps=4000 +# - slurm_config.nodes=8 (global batch stays 64 = nodes x gpus x bs x accum) +# Every other knob below is already the A/B setting. + +job_name: Qwen3-8B_DFlash2_online +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3-8B + + # Step 1: Build input conversations. example_data_config.yaml enables only the + # daring-anteater source (train: 50000) — multi-turn SFT with real assistant + # completions. --full-conversations keeps those completions so answer_only_loss + # has assistant spans to mask. make_dataset.sh writes /scratchspace/data/train.jsonl. + task_0: + script: common/eagle3/make_dataset.sh + args: + - -f modules/Model-Optimizer/examples/dataset/example_data_config.yaml + - --full-conversations + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + + # Step 2: Online DFlash2 training (the script exports the drafter at the end). + # Consumes the conversations built in task_0 (shared via /scratchspace). + task_1: + script: common/specdec/dflash_online_training.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash2.yaml + - model.model_name_or_path=<> + - data.data_path=/scratchspace/data/train.jsonl + - data.chat_template=examples/Qwen/Qwen3-8B/chat_template_train.jinja + - training.output_dir=/scratchspace/dflash2_bs16 + - training.per_device_train_batch_size=1 + - training.num_train_epochs=1 + - training.max_steps=2000 + - training.training_seq_len=4096 + - training.learning_rate=6.0e-4 + - training.warmup_ratio=0.04 + - training.warmup_steps=0 + - training.lr_scheduler_type=linear + - training.save_steps=5000 + - training.logging_steps=100 + - training.disable_tqdm=true + - training.answer_only_loss=true + # Draft backbone — identical to the DFlash / DSpark / Domino arms so the + # only difference between them is the correction head. + - dflash.dflash_block_size=16 + - dflash.dflash_num_anchors=512 + - dflash.dflash_loss_decay_factor=7 + - dflash.dflash_mask_token_id=151669 + - dflash.dflash_self_logit_distillation=false + - dflash.dflash_architecture_config.num_hidden_layers=5 + - dflash.dflash_architecture_config.num_attention_heads=32 + - dflash.dflash_architecture_config.num_key_value_heads=8 + - dflash.dflash_architecture_config.head_dim=128 + - dflash.dflash_architecture_config.intermediate_size=12288 + # DFlash2 knobs (also set in the recipe; repeated here for visibility). + # A draft dim NOT set explicitly falls back to the Qwen3Config default, not + # to the base model's — hence the five dims above are always spelled out. + - dflash.dflash_architecture_config.projector_type=dflash2 + - dflash.dflash_architecture_config.conv_kernel_size=2 + - dflash.dflash_architecture_config.conv_group_size=16 + - dflash.dflash_architecture_config.selector_rank=256 + - dflash.dflash_architecture_config.selector_top_k=16 + - dflash.dflash_selector_loss_alpha=1.0 + # Sliding-window draft attention, matching the published DFlash2 drafters. + - dflash.dflash_swa_window_size=2048 + environment: + - MAX_FINAL_LOSS: "5.0" + - MIN_FINAL_ACC: "0.15" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8