refactor(speculative): draft DSpark and Domino through one AR-validation loop

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -1156,3 +1156,39 @@ def pseudo_speculative_generate(self, input_ids, steps=1):
         # produced on the draft's device, and callers cat both onto the running sequence.
         num_tokens = min(steps, self.dflash_block_size - first)
         return base_token, draft_tokens[:, :num_tokens].to(input_ids.device)
+
+    def _decode_left_to_right(self, input_ids, steps):
+        """``pseudo_speculative_generate`` for a sequential head, drafting the block left to right.
+
+        Block position ``k`` drafts the token after the anchor and the ``k`` tokens drafted before
+        it (``_label_shift = 1``): the argmax of its backbone logits plus the head's correction.
+        """
+        base_token, hidden_states = self._generation_base(input_ids)
+        if steps < 1:
+            return base_token, None
+        draft_hidden = self._draft_block(base_token, hidden_states)
+        device = draft_hidden.device
+        # Under device_map="auto" a tied lm_head sits with the embedding table on the first GPU,
+        # not the draft's, and its output stays there.
+        backbone_logits = self._base_model_lm_head(draft_hidden).to(device)
+
+        prev_token = base_token.squeeze(-1).to(device)
+        state = None
+        draft_tokens = []
+        for k in range(min(steps, self.dflash_block_size)):
+            correction, state = self._correction_step(k, prev_token, draft_hidden[:, k, :], state)
+            logits = backbone_logits[:, k, :]
+            if correction is not None:
+                logits = logits + correction
+            prev_token = logits.argmax(dim=-1)
+            draft_tokens.append(prev_token)
+        return base_token, torch.stack(draft_tokens, dim=1).to(input_ids.device)
+
+    def _correction_step(self, k, prev_token, hidden, state):
+        """A sequential head's correction to block position ``k``'s logits, and its new ``state``.
+
+        ``prev_token`` [B] is the token decoded at ``k - 1`` (the anchor for ``k = 0``), ``hidden``
+        [B, H] the draft's hidden state at ``k``, and ``state`` the head's carry, None at ``k = 0``.
+        The correction is [B, vocab], or None to leave the backbone's logits as they are.
+        """
+        raise NotImplementedError

modelopt/torch/speculative/plugins/hf_domino.py

@@ -205,29 +205,17 @@ def pseudo_speculative_generate(self, input_ids, steps=1):
         before it. From ``pure_draft_prefix_len`` on, its backbone logits get the correction
         from the GRU state over exactly those tokens.
         """
-        base_token, hidden_states = self._generation_base(input_ids)
-        if steps < 1:
-            return base_token, None
-        draft_hidden = self._draft_block(base_token, hidden_states)
-        device = draft_hidden.device
-        # Under device_map="auto" the embedding table and a tied lm_head sit on the first GPU,
-        # not the draft's, and their outputs stay there.
-        backbone_logits = self._base_model_lm_head(draft_hidden).to(device)  # [B, block_size, V]
+        return self._decode_left_to_right(input_ids, steps)
 
+    def _correction_step(self, k, prev_token, hidden, state):
+        """The GRU's correction to block position ``k``; none before ``pure_draft_prefix_len``."""
         m = self.dflash_module
-        prev_token = base_token.squeeze(-1).to(device)
-        gru_state = None
-        draft_tokens = []
-        for k in range(min(steps, self.dflash_block_size)):
-            prev_emb = self._base_model_embeddings(prev_token).to(device)
-            gru_out, gru_state = m.prefix_gru(prev_emb.unsqueeze(1), gru_state)
-            logits = backbone_logits[:, k, :]
-            if k >= m.pure_draft_prefix_len:
-                correction = m.embed_proj(torch.cat([draft_hidden[:, k, :], gru_out[:, 0, :]], -1))
-                logits = logits + correction
-            prev_token = logits.argmax(dim=-1)
-            draft_tokens.append(prev_token)
-        return base_token, torch.stack(draft_tokens, dim=1).to(input_ids.device)
+        # Under device_map="auto" the embedding table sits on the first GPU, not the draft's.
+        prev_emb = self._base_model_embeddings(prev_token).to(hidden.device)
+        gru_out, state = m.prefix_gru(prev_emb.unsqueeze(1), state)
+        if k < m.pure_draft_prefix_len:
+            return None, state
+        return m.embed_proj(torch.cat([hidden, gru_out[:, 0, :]], -1)), state
 
 
 class DominoLambdaCallback(TrainerCallback):

modelopt/torch/speculative/plugins/hf_dspark.py

@@ -193,23 +193,8 @@ def pseudo_speculative_generate(self, input_ids, steps=1):
         (the anchor for k=0). Without this override DSpark would fall back to the
         DFlash backbone-only generate, under-reporting acceptance length.
         """
-        base_token, hidden_states = self._generation_base(input_ids)
-        if steps < 1:
-            return base_token, None
-        draft_hidden = self._draft_block(base_token, hidden_states)
-        # Under device_map="auto" a tied lm_head sits with the embedding table, not the draft.
-        backbone_logits = self._base_model_lm_head(draft_hidden).to(draft_hidden.device)
-
-        # Autoregressive Markov sampling over the block.
-        m = self.dflash_module
-        num_tokens = min(steps, self.dflash_block_size)
-        # anchor precedes block position 0
-        prev_token = base_token.squeeze(-1).to(draft_hidden.device)
-        state = None
-        draft_tokens = []
-        for k in range(num_tokens):
-            bias, state = m.markov_step(prev_token, draft_hidden[:, k, :], state)
-            tok = (backbone_logits[:, k, :] + bias).argmax(dim=-1)
-            draft_tokens.append(tok)
-            prev_token = tok
-        return base_token, torch.stack(draft_tokens, dim=1).to(input_ids.device)
+        return self._decode_left_to_right(input_ids, steps)
+
+    def _correction_step(self, k, prev_token, hidden, state):
+        """Block position ``k``'s Markov bias; only the rnn head carries ``state``."""
+        return self.dflash_module.markov_step(prev_token, hidden, state)

tests/unit/torch/speculative/plugins/test_hf_dspark.py

@@ -269,6 +269,40 @@ def spy(*args, **kwargs):
         assert draft_tokens.shape == (1, 3)
 
 
+class TestDSparkGenerate:
+    """AR validation drafts with the head that is trained and exported."""
+
+    @pytest.mark.parametrize("head_type", HEAD_TYPES)
+    def test_drafts_are_the_training_head_argmax(self, head_type):
+        """Greedy drafts equal the training head's argmax when it is teacher-forced on them."""
+        model = get_tiny_llama(num_hidden_layers=4)
+        mtsp.convert(model, [("dflash", _get_dspark_config(head_type=head_type))])
+        model.float().eval()
+        torch.manual_seed(0)
+        with torch.no_grad():
+            # Large enough to overturn the backbone's argmax, so a draft that skips the bias or
+            # feeds the head the wrong prefix, hidden state or recurrent state disagrees.
+            for name, p in model.dflash_module.named_parameters():
+                if name.startswith(("markov_", "gate_proj.", "joint_proj.")):
+                    p.normal_()
+        input_ids = torch.randint(1, model.dflash_config.vocab_size, (2, SEQ_LEN))
+
+        base_token, drafts = model.pseudo_speculative_generate(input_ids, steps=BLOCK_SIZE)
+
+        with torch.no_grad():
+            _, hidden_states = model._generation_base(input_ids)
+            batch = DraftBatch(
+                torch.cat([base_token, drafts], dim=1),  # the anchor, then the drafts
+                loss_mask=None,
+                anchors=torch.zeros(2, 1, dtype=torch.long),
+                keep=None,
+                base_outputs=None,
+                hidden=model._draft_block(base_token, hidden_states),
+            )
+            final_logits = model._draft_logits(batch).final  # [B, 1 block, bs, V]
+        assert torch.equal(final_logits[:, 0].argmax(dim=-1), drafts)
+
+
 class TestDSparkExporter:
     """Test vanilla serving exports and the native gated/recurrent formats."""