Support train Domino draft model?

[lengrongfu]

https://github.com/jianuo-huang/Domino

[yubofredwang]

yes, we do plan to support, but might take some time

[SeaTuKeMa]

Hi maintainers, I’m interested in implementing Domino training support.

Proposed approach:

1. Reuse the existing DFlash backbone and training infrastructure from feat: DFlash (block-parallel) draft model training #64.
2. Add optional Domino configuration and its GRU-based causal correction head.
3. Add a DominoTrainer extending DFlashTrainer.
4. Train with combined base DFlash loss and Domino-corrected loss.
5. Add a Qwen3-8B example configuration.

Test plan:

• Domino configuration validation
• Correction-head output shapes and causality
• Base/corrected loss interpolation
• Shifted-label and sequence-boundary handling
• Existing DFlash behavior remains unchanged
• Small end-to-end training smoke test

Would like to listen whats your design plan as well, could follow your prefered direction if your team already have a existing design
cc @yubofredwang

[yubofredwang]

@ZYInSea I feel that is a valid plan. Feel free to implement end to end. I can help with the validation as well with the GPU resources I have

[SeaTuKeMa]

WIP in PR#118,did initial test on H100, will try run on multiple gpu settings and provide experiment results

[SeaTuKeMa]

WIP in PR#118,did initial test on H100, will try run on multiple gpu settings and provide experiment results

Ready for review @yubofredwang