[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.

[slepshep106]

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?

Thanks for filing this — I ran into the same need and came across your add_vl_support branch while investigating. Great work!

Two questions:

1. Qwen3.5 image-text DFlash — is configs/qwen3.5-35b-a3b-dflash-vlm-8layer.json expected to work end-to-end today? The status doc lists a live multimodal capture run and an end-to-end training run as "not yet verified (needs GPU/NPU)". We're targeting Qwen3.5-4B and Qwen3.5-35B-A3B on OCR data.

2. DSpark — modality="multimodal" is currently registered only in dflash/providers.py; algorithms/dspark/ is untouched on the branch. Is multimodal DSpark on your roadmap, or would you be open to a PR?

Thanks again for pushing this forward!

[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?

Thanks for filing this — I ran into the same need and came across your add_vl_support branch while investigating. Great work!

Two questions:

1. Qwen3.5 image-text DFlash — is configs/qwen3.5-35b-a3b-dflash-vlm-8layer.json expected to work end-to-end today? The status doc lists a live multimodal capture run and an end-to-end training run as "not yet verified (needs GPU/NPU)". We're targeting Qwen3.5-4B and Qwen3.5-35B-A3B on OCR data.
2. DSpark — modality="multimodal" is currently registered only in dflash/providers.py; algorithms/dspark/ is untouched on the branch. Is multimodal DSpark on your roadmap, or would you be open to a PR?

Thanks again for pushing this forward!

Hi @slepshep106, thanks for the interest!

Q1: Implemented, but not yet validated on GPU/NPU — the branch has the full path (multimodal FeatureContract + VLM input adapter, a regenerated v0.5.14 spec-capture patch with a position_ids mRoPE artifact, collator + 3D mRoPE positions in the DFlash training forward). We're doing the first live bring-up on Ascend NPU / Qwen3.5-4B right now and expect to finish the initial evaluation and open the PR within this week; main open risk is the sglang server side (ViT + aux capture). If you try it: patch sglang 0.5.14 with the branch's patch, set input_modality: multimodal, one image per sample, keep client/server transformers versions aligned, and use target-regenerated data (per #585, accept length depends heavily on it). Status doc: docs/advanced_features/vlm_dflash.md.

Q2: Once our end-to-end DFlash work lands upstream, we'll start the DSpark adaptation — and the path is already prepared for it: DSpark shares _forward_draft_blocks with DFlash (already takes position_ids), and the capture patch already provides the last_hidden artifact its L1/confidence losses need, so the adaptation is mainly registering the contract + provider + collator in algorithms/dspark/providers.py, mirroring dflash.

[slepshep106]

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?

Thanks for filing this — I ran into the same need and came across your add_vl_support branch while investigating. Great work!
Two questions:

1. Qwen3.5 image-text DFlash — is configs/qwen3.5-35b-a3b-dflash-vlm-8layer.json expected to work end-to-end today? The status doc lists a live multimodal capture run and an end-to-end training run as "not yet verified (needs GPU/NPU)". We're targeting Qwen3.5-4B and Qwen3.5-35B-A3B on OCR data.
2. DSpark — modality="multimodal" is currently registered only in dflash/providers.py; algorithms/dspark/ is untouched on the branch. Is multimodal DSpark on your roadmap, or would you be open to a PR?

Thanks again for pushing this forward!

Hi @slepshep106, thanks for the interest!

Q1: Implemented, but not yet validated on GPU/NPU — the branch has the full path (multimodal FeatureContract + VLM input adapter, a regenerated v0.5.14 spec-capture patch with a position_ids mRoPE artifact, collator + 3D mRoPE positions in the DFlash training forward). We're doing the first live bring-up on Ascend NPU / Qwen3.5-4B right now and expect to finish the initial evaluation and open the PR within this week; main open risk is the sglang server side (ViT + aux capture). If you try it: patch sglang 0.5.14 with the branch's patch, set input_modality: multimodal, one image per sample, keep client/server transformers versions aligned, and use target-regenerated data (per #585, accept length depends heavily on it). Status doc: docs/advanced_features/vlm_dflash.md.

Q2: Once our end-to-end DFlash work lands upstream, we'll start the DSpark adaptation — and the path is already prepared for it: DSpark shares _forward_draft_blocks with DFlash (already takes position_ids), and the capture patch already provides the last_hidden artifact its L1/confidence losses need, so the adaptation is mainly registering the contract + provider + collator in algorithms/dspark/providers.py, mirroring dflash.

Thanks, very helpful.

We'll start bring-up on GPU with Qwen3.5-4B for OCR this week following your notes. We'll report back on the sglang server side specifically, and file issues against the branch if we hit anything.