[h-guo18]

What does this PR do?

Type of change: Refactor. No behavior change: the drafted tokens are bitwise identical.

This PR is stacked on #2990 (branch haoguo/dflash-shared-forward, head 28f4f45), whose branch is its base; merge it after that one. The stack below it, on main 0638a27, is #2757 ← #2762 ← #2989 [1/2] (branch haoguo/dflash-loss-context) ← #2990 [2/2]. Three other PRs also sit on #2990 and are independent of this one: #2991 (DSpark and Domino follow dflash_loss_objective), #2992 (logging draft_metrics) and #2993 (loss terms in config). This branch merges with each of them without conflicts, and none of them touches the generation path.

#2762 gave Domino an AR-validation decode modeled on DSpark's, so pseudo_speculative_generate in hf_dspark.py and hf_domino.py carried the same loop over the block:

• the base forward and the draft block;
• the move of the lm_head's logits onto the draft's device (fix(speculative): DSpark and Domino AR validation and label alignment #2762's device_map="auto" fix, in both copies);
• the previous token and the head's state;
• the cap at the block size, the stacking, and the move back to the caller's device.

Only the step at each block position differed. A fix to the loop had to be made twice, and a third variant with a sequential head would have copied it again.

• One loop. HFDFlashModel._decode_left_to_right(input_ids, steps) holds the loop and the device handling. Block position k drafts the token after the anchor and the k tokens drafted before it (_label_shift = 1).
• One step per variant. A variant supplies _correction_step(k, prev_token, hidden, state). It returns the correction to position k's backbone logits, or None to leave them as they are, together with the head's new state.
  • DSpark: markov_step. Only the rnn head carries state.
  • Domino: steps the GRU over the previous token's embedding, which it still moves onto the draft's device. From pure_draft_prefix_len on, it returns embed_proj's correction.
• Unchanged. 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.
• New test. DSpark's decode had only shape tests. TestDSparkGenerate mirrors fix(speculative): DSpark and Domino AR validation and label alignment #2762's TestDominoGenerate for the vanilla, gated and rnn heads: with head weights large enough to overturn the backbone, the greedy drafts must equal the argmax of _draft_logits teacher-forced on them. It lands first and passes on [2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms #2990's head.

Usage

No API or config change. A variant with a sequential head now plugs in like this:

class HFMyDraftModel(HFDFlashModel):
    _label_shift = 1

    @torch.no_grad()
    def pseudo_speculative_generate(self, input_ids, steps=1):
        return self._decode_left_to_right(input_ids, steps)

    def _correction_step(self, k, prev_token, hidden, state):
        # The correction to block position k's backbone logits (or None), and the new state.
        return self.dflash_module.my_head_step(prev_token, hidden, state)

Testing

• The DFlash-family unit tests (tests/unit/torch/speculative/plugins/test_hf_{dflash,dflash2,dspark,domino,lilicorr}.py) plus tests/unit/torch/export/test_hf_spec_rope_export.py: 348 passed on transformers 5.5.4, 5.12.1 and 5.18.0. That is [2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms #2990's 345 tests in these files plus the 3 new cases, each under 0.2 s on 4 CPUs. 5.5.4 and 5.18.0 fall within CI's tf_min and tf_latest pins. fix(speculative): DSpark and Domino AR validation and label alignment #2762's TestDominoGenerate (pure_draft_prefix_len 0 and 1) is unchanged and passes.

• tests/gpu/torch/speculative/plugins on one GPU: 20 passed, 1 skipped (it needs 2 GPUs). This includes fix(speculative): DSpark and Domino AR validation and label alignment #2762's TestLmHeadOffTheDraftDevice for DFlash, DSpark and Domino, unchanged.

• Bitwise A/B against [2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms #2990's head. It compares loss, every gradient, train_acc, draft_metrics and the drafted tokens. All 30 cases are identical in fp32 and bf16 (60/60).

  • The cases cover DFlash, DFlash2, LiLiCorr, DSpark and Domino, offline and online, including zero-loss steps.
  • The online DSpark (vanilla, gated and rnn heads) and Domino (pure_draft_prefix_len 1 and 0) cases draft with steps 0, 1, 2, 4 and 6, the last one past the block size.
  • They draft at init scale, and again from 9- and 13-token prefixes with head weights large enough to decide the drafts. That gives 26 generated tensors per run.

• Sharded generation on one GPU, with the CPU standing in for the first GPU (accelerate's dispatch_model). DFlash, the three DSpark heads and Domino with pure_draft_prefix_len 1 and 0 were drafted in fp32 and bf16, with steps 0, 1, 2, 4 and 6 from three prefix lengths, under two device maps:

  • embed_tokens and lm_head on the CPU, the rest on the GPU, as in TestLmHeadOffTheDraftDevice;
  • the same, plus base layers 0-1 on the CPU, with target_layer_ids [0, 3] so the target hidden states come off both devices.

  All 24 runs are bitwise identical to [2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms #2990's head, and the base token and drafts come back on the caller's device.

• Mutation check. Each edit below fails at least one of the unit tests above, the GPU test or the A/B:

  • in the shared loop: skipping the correction, dropping the head's state between positions, never updating the previous token, reading the next position's hidden state or backbone logits, capping one position short, and leaving out the lm_head or return-device move;
  • in the steps: DSpark's step dropping the rnn state (of the unit tests, only the new one catches it), Domino's prefix boundary off by one, and Domino's GRU skipping the prefix positions.

  Removing Domino's move of the token embedding passes everything. Under accelerate's dispatch, the GRU's own hook moves its input. The move stays as fix(speculative): DSpark and Domino AR validation and label alignment #2762 wrote it.

• FSDP2 probe (2 CPU ranks on gloo, the draft as its own fully_shard unit, one rank on the zero-loss branch). Training is untouched, and DFlash under CE and KD finishes with the same losses as on [2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms #2990's head. Domino fails under this probe with a DTensor/Tensor mix, identically on [2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms #2990's head; that predates this PR.

• pre-commit passes on the changed files.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅ Only private methods change; configs, checkpoints and pseudo_speculative_generate's outputs are unchanged.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅ TestDSparkGenerate.
• Did you update Changelog?: N/A Refactor with no user-visible change.
• Did you get Claude approval on this PR?: ❌ Not run yet.

Additional Information

Source changes: +50 / −41 lines in hf_dflash.py, hf_dspark.py and hf_domino.py. Two commits: the DSpark test, which passes on the parent, then the refactor.

Each variant opts in with a one-line pseudo_speculative_generate override rather than by defining _correction_step alone. Switching on whether a variant defines the step would need a None sentinel on HFDFlashModel and a branch in DFlash's decode. Making the loop the family's only decode, with no correction for DFlash, would turn DFlash's parallel decode into a loop and move DFlash2 and LiLiCorr onto it. Both are larger than this refactor. The base _correction_step raises NotImplementedError, so a variant that opts in without a step fails on its first draft.

Merge order: #2757, #2762, #2989, #2990, then this PR.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}