[Feature] Support image-text multimodal input for DFlash training, similar to Eagle3

[curnane-lab]

Checklist

• 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/SpecForge/discussions/new/choose Otherwise, it will be closed.
• 2. Please use English, otherwise it will be closed.

Motivation

Checklist

• 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/SpecForge/discussions/new/choose Otherwise, it will be closed.
• 2. Please use English, otherwise it will be closed.

Motivation

Currently, the DFlash training pipeline in scripts/train_dflash.py is designed for text-only data. However, many real-world speculative decoding scenarios involve multimodal content (e.g., image + text). Eagle3 has demonstrated support for multimodal inputs, and it would be valuable for SpecForge's DFlash implementation to support image-text inputs as well.

Specifically, we would like to train DFlash draft models for vision-language models such as Qwen3.5/Qwen3.6, where the target model processes both images and text, and the draft model accelerates token generation by predicting future text tokens conditioned on the multimodal context.

Related resources

• Eagle3 (supports multimodal speculative decoding): https://github.com/NVlabs/Eagle
• Qwen3.5 model family (potential target models)
• Existing flexible embedding key in SpecForge: --embedding-key (e.g., model.language_model.embed_tokens.weight)

Proposed scope

1. Data pipeline: Support loading image-text paired data in build_eagle3_dataset or a dedicated DFlash multimodal dataset builder.
2. Model inputs: Allow the target model to accept pixel values / image embeddings alongside input_ids and attention_mask.
3. Hidden-state capture: Correctly capture hidden states from the target multimodal language model at the specified target_layer_ids, taking into account vision-language fusion.
4. Draft model alignment: Ensure the DFlash draft model can consume the multimodal context and produce block-level draft predictions.
5. Example script & config: Provide an example training script and a reference config for a multimodal DFlash draft model.

Related resources

No response

[curnane-lab]

Implementation Plan

I have a working prototype at my private repo and propose to split it into 4 sequential PRs:

1. [refactor] Decouple DFlash dataset builder from Eagle3 — paves the way
   for VLM data pipeline without breaking text-only training.
2. [feat] Add build_dflash_vlm_dataset for image-text pairs.
3. [feat] Extend DFlashTargetModel / DFlashDraftModel to accept
   pixel_values / image_grid_thw (optional, backward-compatible).
4. [example] Qwen2.5-VL-7B DFlash online training example + benchmark.

Targets:

• Qwen2.5-VL-7B / Qwen3.5-series
• Reuse Eagle3's QwenVLOnlineEagle3Model interface for consistency
• Zero regression on pure-text DFlash

Could a maintainer (@FlamingoPg / @sleepcoo/ @jiahang01 @jiapingW ) confirm this split is acceptable
before I open the first PR?

[jiapingW]

Hi，Eagle3 and DFlash should share a set of VLM and LLM data input interfaces. However, the acquisition of hidden states and other information needs to be adjusted according to the scenario. Text input and VLM input also use the same interface, which requires some encapsulation.