refactor(speculative): one home each for the DFlash family's shared rules

modelopt/torch/speculative/dflash/conversion.py

@@ -48,6 +48,16 @@
 # registry so its wrapper (HFLiLiCorrModel) does not overwrite HFDFlashModel.
 LiLiCorrDMRegistry = _DMRegistryCls(prefix="LiLiCorr")
 
+# dflash_architecture_config.projector_type -> the registry holding that variant's wrapper.
+_REGISTRY_BY_PROJECTOR_TYPE = {
+    None: DFlashDMRegistry,
+    "dflash": DFlashDMRegistry,
+    "domino": DominoDMRegistry,
+    "dspark": DSparkDMRegistry,
+    "dflash2": DFlash2DMRegistry,
+    "lilicorr": LiLiCorrDMRegistry,
+}
+
 
 def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertReturnType:
     """Convert the model to a DFlash (or Domino) model as per `config`."""
@@ -59,22 +69,13 @@ def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertRe
     custom_config = config.dflash_architecture_config
     config.dflash_architecture_config = {**default_dflash_config, **custom_config}
 
-    # Route to the Domino registry when the architecture asks for the Domino head.
     projector_type = config.dflash_architecture_config.get("projector_type")
-    if projector_type == "domino":
-        registry = DominoDMRegistry
-    elif projector_type == "dspark":
-        registry = DSparkDMRegistry
-    elif projector_type == "dflash2":
-        registry = DFlash2DMRegistry
-    elif projector_type == "lilicorr":
-        registry = LiLiCorrDMRegistry
-    elif projector_type in (None, "dflash"):
-        registry = DFlashDMRegistry
-    else:
+    registry = _REGISTRY_BY_PROJECTOR_TYPE.get(projector_type)
+    if registry is None:
+        supported = [name for name in _REGISTRY_BY_PROJECTOR_TYPE if name is not None]
         raise ValueError(
             f"Unsupported dflash_architecture_config.projector_type: {projector_type!r}. "
-            "Expected 'dflash' (default), 'domino', 'dspark', 'dflash2' or 'lilicorr'."
+            f"Expected one of {supported} ('dflash' is the default)."
         )
 
     registry.register_subclass(type(model), "base_model_class")

modelopt/torch/speculative/plugins/hf_dflash2.py

@@ -39,7 +39,7 @@
 from transformers import PreTrainedModel
 
 from ..dflash.conversion import DFlash2DMRegistry
-from .hf_dflash import HFDFlashModel
+from .hf_dflash import HFDFlashModel, _block_targets, _block_tokens
 from .modeling_dflash2 import DFlash2Module
 
 __all__ = ["HFDFlash2Model"]
@@ -72,10 +72,10 @@ def modify(self, config):
                 "candidate selector's rank/top-k)."
             )
         super().modify(config)
-        self.dflash_selector_loss_alpha = getattr(config, "dflash_selector_loss_alpha", 1.0)
-        self.dflash_lk_loss_type = getattr(config, "dflash_lk_loss_type", "ce")
-        self.dflash_lk_ce_scale = getattr(config, "dflash_lk_ce_scale", 1.0)
-        self.dflash_lk_ce_decay = getattr(config, "dflash_lk_ce_decay", 1.0)
+        self.dflash_selector_loss_alpha = config.dflash_selector_loss_alpha
+        self.dflash_lk_loss_type = config.dflash_lk_loss_type
+        self.dflash_lk_ce_scale = config.dflash_lk_ce_scale
+        self.dflash_lk_ce_decay = config.dflash_lk_ce_decay
         if self.dflash_lk_loss_type != "ce" and self.dflash_self_logit_distillation:
             raise ValueError(
                 f"dflash_lk_loss_type={self.dflash_lk_loss_type!r} needs the draft's "
@@ -231,45 +231,35 @@ def _compute_loss(
         if self.dflash_selector_loss_alpha <= 0 or draft_hidden is None:
             return loss, accuracy
 
-        bsz, seq_len = input_ids.shape
+        bsz = input_ids.shape[0]
         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
+
+        # The backbone's supervision mask, without its position weighting (decay/D-PACE):
+        # that shapes *where* the backbone spends capacity, while the selector should learn
+        # every position's transition equally.
+        targets = _block_targets(
+            input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, self._label_shift
         )
 
         # Teacher-forced predecessor of block position k is the real token at anchor+k-1.
-        # Only offsets 1..block_size-1 are supervised (weight_mask zeroes slot 0), so the
+        # Only offsets 1..block_size-1 are supervised (the mask zeroes slot 0), so the
         # first supervised position, k=1, has the anchor's own token as its predecessor.
         # Slot 0's entry here resolves to anchor-1 and never contributes.
         #
         # This is the train/serve contract: CandidateSelector.greedy_path seeds its walk
         # from the anchor token, so ITS position 0 corresponds to block offset 1. A caller
         # that hands greedy_path the full 0..block_size-1 candidate set is off by one.
-        predecessor_ids = torch.gather(expanded_ids, 2, (safe_label_indices - 1).clamp(min=0))
+        predecessor_ids = _block_tokens(
+            input_ids, anchor_positions, block_size, self._label_shift - 1
+        )
 
         selector_loss, selector_accuracy, selector_coverage = self._selector_loss(
             logits.reshape(bsz, n_blocks, block_size, -1),
-            target_ids,
+            targets.token_ids,
             draft_hidden.reshape(bsz, n_blocks, block_size, -1),
             predecessor_ids,
-            weight_mask,
+            targets.mask,
         )
         self._selector_metrics = {
             "selector_accuracy": selector_accuracy,

modelopt/torch/speculative/plugins/hf_domino.py

@@ -68,7 +68,14 @@
 from transformers.utils import ModelOutput
 
 from ..dflash.conversion import DominoDMRegistry
-from .hf_dflash import HFDFlashModel
+from .hf_dflash import (
+    HFDFlashModel,
+    _block_position_weights,
+    _block_targets,
+    _block_tokens,
+    _first_predicted_position,
+    _masked_top1,
+)
 from .modeling_dflash import DFlashBaseModelOutput
 from .modeling_domino import DominoModule
 
@@ -101,6 +108,8 @@ class HFDominoModel(HFDFlashModel):
     route to it when ``dflash_architecture_config.projector_type == "domino"``.
     """
 
+    _label_shift = 1
+
     def _build_draft_module(self, dflash_config):
         """Build the Domino draft module (DFlash backbone + GRU correction head)."""
         return DominoModule(dflash_config)
@@ -121,8 +130,8 @@ def modify(self, config):
         # (DFlashConfig carries the two fields); updated each step by
         # DominoLambdaCallback. Defaults keep a single forward (e.g. unit tests)
         # well-defined without a scheduler.
-        self.dflash_lambda_base_start = getattr(config, "dflash_lambda_base_start", 1.0)
-        self.dflash_lambda_base_decay_ratio = getattr(config, "dflash_lambda_base_decay_ratio", 1.0)
+        self.dflash_lambda_base_start = config.dflash_lambda_base_start
+        self.dflash_lambda_base_decay_ratio = config.dflash_lambda_base_decay_ratio
         self._lambda_base = self.dflash_lambda_base_start
         if not getattr(self.dflash_module, "shift_label", True):
             raise NotImplementedError(
@@ -151,18 +160,15 @@ def _apply_domino_head(self, hidden, base_logits, input_ids, anchor_positions, n
         Returns:
             Corrected logits [B, N*block_size, vocab].
         """
-        bsz, seq_len = input_ids.shape
+        bsz = input_ids.shape[0]
         bs = self.dflash_block_size
-        device = input_ids.device
         suffix_start = self.dflash_module.pure_draft_prefix_len
 
         hidden4d = hidden.reshape(bsz, n_blocks, bs, hidden.size(-1))
         base4d = base_logits.reshape(bsz, n_blocks, bs, -1)
 
         # Teacher-forced previous tokens: the real token at anchor+j for j in [0, bs).
-        prev_offsets = torch.arange(bs, device=device).view(1, 1, -1)
-        prev_idx = (anchor_positions.unsqueeze(-1) + prev_offsets).clamp(max=seq_len - 1)
-        prev_ids = torch.gather(input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, prev_idx)
+        prev_ids = _block_tokens(input_ids, anchor_positions, bs, self._label_shift - 1)
 
         block_emb = self._base_model_embeddings(prev_ids)  # [B, N, bs, H]
         gru_in = block_emb.reshape(bsz * n_blocks, bs, block_emb.size(-1))
@@ -187,40 +193,27 @@ def _compute_domino_loss(
         Uses next-token (shift_label) alignment: position k predicts the token at
         anchor+k+1, and position 0 is *not* excluded (unlike base DFlash).
         """
-        bsz, seq_len = input_ids.shape
-        bs = self.dflash_block_size
-        n_blocks = anchor_positions.shape[1]
-        device = input_ids.device
-
-        # shift_label=True: label for block position k is the token at anchor+k+1.
-        label_offsets = torch.arange(1, 1 + bs, device=device).view(1, 1, -1)
-        label_indices = anchor_positions.unsqueeze(-1) + label_offsets
-        valid_label = label_indices < seq_len
-        safe_label_indices = label_indices.clamp(max=seq_len - 1)
-
-        target_ids = torch.gather(
-            input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
+        targets = _block_targets(
+            input_ids,
+            anchor_positions,
+            block_keep_mask,
+            loss_mask,
+            self.dflash_block_size,
+            self._label_shift,
         )
-
-        # Weight mask: valid block * in bounds * loss_mask. No pos-0 exclusion.
-        weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, bs).float()
-        weight_mask = weight_mask * valid_label.float()
-        orig_loss_mask = torch.gather(
-            loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
+        binary_eval_mask = targets.mask.view(-1)
+        # dflash_loss_objective is not wired for Domino yet: it always decays.
+        weight_mask = _block_position_weights(
+            targets.mask,
+            "decay",
+            _first_predicted_position(self._label_shift),
+            decay_factor=self.dflash_loss_decay_factor,
+            dpace_alpha=self.dflash_dpace_alpha,
         )
-        weight_mask = weight_mask * orig_loss_mask
-
-        binary_eval_mask = weight_mask.view(-1)
-
-        # Loss decay: exp(-k/gamma) so the first prediction (k=0) gets weight 1.0.
-        if self.dflash_loss_decay_factor > 0:
-            k = torch.arange(bs, device=device).view(1, 1, -1)
-            decay = torch.exp(-k.clamp(min=0).float() / self.dflash_loss_decay_factor)
-            weight_mask = weight_mask * decay
 
         flat_final = final_logits.reshape(-1, final_logits.size(-1))
         flat_base = base_logits.reshape(-1, base_logits.size(-1))
-        flat_targets = target_ids.reshape(-1)
+        flat_targets = targets.token_ids.reshape(-1)
         flat_weights = weight_mask.reshape(-1)
         valid_count = flat_weights.sum() + 1e-6
 
@@ -236,15 +229,10 @@ def _compute_domino_loss(
             loss = (1.0 - lambda_base) * final_loss + lambda_base * base_loss
 
             with torch.no_grad():
-                eval_count = binary_eval_mask.sum() + 1e-6
-                final_correct = (flat_final.argmax(dim=-1) == flat_targets) & (
-                    binary_eval_mask > 0.5
-                )
-                base_correct = (flat_base.argmax(dim=-1) == flat_targets) & (binary_eval_mask > 0.5)
-                accuracy = (final_correct.sum().float() / eval_count).item()
-                base_accuracy = (base_correct.sum().float() / eval_count).item()
+                accuracy = _masked_top1(flat_final, flat_targets, binary_eval_mask).item()
+                base_accuracy = _masked_top1(flat_base, flat_targets, binary_eval_mask).item()
         else:
-            loss = flat_final.sum() * 0.0
+            loss = self._draft_zero_loss(flat_final.device, through=(flat_final,))
             final_loss = loss
             base_loss = loss
             accuracy = 0.0
@@ -323,9 +311,7 @@ def forward(
                     attention_mask=attention_mask,
                     output_hidden_states=True,
                 )
-            offset = 1
-            selected = [base_out.hidden_states[lid + offset] for lid in self.target_layer_ids]
-            target_hidden = torch.cat(selected, dim=-1)  # [B, seq, num_layers * H]
+            target_hidden = self._target_hidden(base_out.hidden_states)
 
         # 2. Build loss mask (same convention as DFlash).
         if labels is not None:
@@ -344,9 +330,7 @@ def forward(
         n_blocks = anchor_positions.shape[1]
 
         if n_blocks == 0 or not block_keep_mask.any():
-            # Zero loss that still flows through all draft params for DDP sync.
-            dummy = sum(p.sum() for p in self.dflash_module.parameters()) * 0.0
-            return ModelOutput(loss=dummy, logits=None, train_acc=[[0.0]])
+            return ModelOutput(loss=self._draft_zero_loss(device), logits=None, train_acc=[[0.0]])
 
         # 4. Build draft inputs (inherited helpers).
         noise_embedding = self._build_noise_embedding(