From a26375590183e5857f4cab05181790529b01cd29 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:02:06 +0000 Subject: [PATCH 1/2] test(speculative): pin DSpark's AR-validation drafts to its training head DSpark's pseudo_speculative_generate drafts the block left to right, one markov_step per position, while training applies the Markov head to the whole block at once through _draft_logits. Only shape tests covered the decode, so nothing checked that the two agree, and the next commit moves the decode loop into HFDFlashModel. The test mirrors Domino's TestDominoGenerate: with head weights large enough to overturn the backbone, the greedy drafts must equal the argmax of _draft_logits teacher-forced on the anchor and those drafts. It runs for the vanilla, gated and rnn heads. Verified: passes on the parent commit for all three heads, in under 0.2 s per case on 4 CPUs. On the next commit it is the only unit test that fails when DSpark's step drops the rnn head's state. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../speculative/plugins/test_hf_dspark.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py index 2e2c541a090..b60d754c89b 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py +++ b/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.""" From 5a4299d3a1a48f97b4498eef65fc247cf42e0566 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:02:38 +0000 Subject: [PATCH 2/2] refactor(speculative): draft DSpark and Domino through one AR-validation loop DSpark's and Domino's pseudo_speculative_generate each carried the same left-to-right loop over the block. Each copy ran the base forward and the draft block, and moved the lm_head's logits onto the draft's device (the #2762 fix, present in both copies). Each also kept the previous token and the head's state, capped the drafts at the block size, stacked them and moved them back to the caller's device. Only the step at each position differed. The loop is now HFDFlashModel._decode_left_to_right. A variant supplies _correction_step(k, prev_token, hidden, state), which returns the correction to block position k's backbone logits (None for none) and the head's new state: - DSpark's step is markov_step. - Domino's step runs the GRU over the previous token's embedding, which it still moves onto the draft's device, and adds embed_proj's correction from pure_draft_prefix_len on. Both pseudo_speculative_generate overrides keep their docstrings and call the shared loop. DFlash's parallel decode, which DFlash2 and LiLiCorr also use, is unchanged. Verified against the parent commit: - Bitwise A/B: the drafted tokens of the online DSpark (vanilla, gated and rnn heads) and Domino (pure_draft_prefix_len 1 and 0) cases are identical in fp32 and bf16, at init scale and with head weights large enough to decide the drafts. Loss, gradients and metrics of all 30 cases, offline and online, are identical too. - Sharded on one GPU with the CPU as the first device: accelerate dispatch with embed_tokens and lm_head on the CPU, and a split that also puts base layers 0-1 there. The drafts are identical for DFlash, DSpark and Domino in fp32 and bf16. - Unit tests in test_hf_{dflash,dflash2,dspark,domino,lilicorr}.py and test_hf_spec_rope_export.py: 348 passed on transformers 5.5.4, 5.12.1 and 5.18.0. GPU plugin tests: 20 passed, 1 skipped (needs 2 GPUs). - Mutation check: 11 of 12 edits to the loop or the steps fail a unit test, the lm_head GPU test or the A/B. Dropping Domino's embedding move passes everything, because accelerate's hook on the GRU moves its input. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../torch/speculative/plugins/hf_dflash.py | 36 +++++++++++++++++++ .../torch/speculative/plugins/hf_domino.py | 30 +++++----------- .../torch/speculative/plugins/hf_dspark.py | 25 +++---------- 3 files changed, 50 insertions(+), 41 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index a0c2cc34bb1..b686478232c 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/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 diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index 0a4bc9fdbd4..661810967a3 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/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): diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index 69b5fdc3f1f..e93ea5b09d4 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/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)