[Feature] Add VP drafter training mode for DFlash

[catnanami]

Motivation

This PR adds training support for the VP-Drafter used in D2SD (Dual Diffusion Draft Speculative Decoding). D2SD extends DFlash by using a first DFlash draft to estimate likely rejection boundaries, then training a second variable-prefix drafter to re-anchor at selected prefixes and generate alternative continuations.

The key training requirement is different from standard DFlash: the drafter must learn from variable-length visible prefixes instead of always seeing only the anchor token followed by masks. This PR implements that behavior as a DFlash training-mode branch, so the resulting model still uses the same DFlashDraftModel architecture and config format.

References:

• Project: https://github.com/catnanami/D2-SD
• Paper: https://arxiv.org/abs/2606.04446

Modifications

• Added a vp_drafter training mode to OnlineDFlashModel.
• The vp_drafter mode samples a variable visible prefix length per block.
• Prefix tokens are fed as real token embeddings; suffix positions remain masked and contribute to loss.
• Added exponential loss decay from the first masked suffix position, matching the VP-Drafter training recipe.
• Wired scripts/train_dflash.py to read dflash_config.training_mode / dflash_config.loss_type when --loss-type is not explicitly provided.
• Added prefix_weight_base support for variable-prefix sampling.
• Added a Qwen3-8B VP-Drafter config and example script:
  • configs/qwen3-8b-dta.json
  • examples/run_qwen3_8b_dta_online.sh
• Added formula-level unit coverage for VP-Drafter loss masking and decay behavior.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[gemini-code-assist]

Code Review

This pull request introduces the vp_drafter training objective (variable-prefix drafter) to the DFlash framework, allowing for training with variable visible prefixes. It adds a configuration file for Qwen3-8B, an online training script, and updates the core DFlash model and training script to support prefix length sampling, VP noise embedding generation, and corresponding loss calculations. Unit tests are also added to validate the new loss implementation. The review feedback highlights three key improvements: replacing torch.distributions.Categorical with torch.multinomial to prevent graph breaks under torch.compile, handling potential None values for prefix_weight_base to avoid a TypeError, and using .reshape() instead of .view() on noise_ids to safely handle non-contiguous tensors.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[catnanami]

Hi @claude @jiapingW @wenqf11, could you please review this PR?

[jiapingW]

Great work! We'll check it soon.