[Feature] Multimodal (image+text) DFlash training via server-side capture (Qwen3.5)

[curnane-lab]

Motivation

Image-text DFlash training is a long-standing request (#583, #461) and a high-priority item on the 2026 roadmap (#374: "General VLM training"). The pre-#678 HF-backend stack had community VLM prototypes (#585, #495, #505); after the unified-runtime consolidation (#678), training runs exclusively on server-side capture, and none of those prototypes apply anymore. The serving side already runs DFlash speculative decoding on both text and image-text inputs - the training side (SpecForge DFlash draft training) is text-only. This PR closes that training-side gap end-to-end on the current architecture.

Design: plain-rope draft (mRoPE-free)

The draft consumes visual information exclusively through the captured target hidden states - the target forward (ViT fusion + the target's own mRoPE) already encodes vision into the context features. The draft's rope is only an internal distance metric for its own attention, decoupled from the target's position scheme, so the target's (3, N) mRoPE positions carry no signal for the draft.

The draft therefore keeps the plain 1D rope convention - byte-identical to the text path - and reuses the official draft geometry (configs/qwen3.5-4b-dflash.json, matching z-lab/Qwen3.5-4B-DFlash). Two hard wins:

• Train/serve consistency for free: the trained draft keeps the same plain-rope convention as the released official draft, so artifacts serve directly on the stock serving path - no engine-side fork patches, no mRoPE train/serve parity audits.
• Draft side is fully isomorphic to text-only: all VLM complexity stays inside "how input enters the target"; multimodal batches flow through the identical training forward.

Modifications

• Contracts/providers (algorithms/dflash/providers.py, common/providers.py): DFlash registers a STREAMING multimodal FeatureContract ({input_ids, loss_mask, hidden_states} - the same three tensors as text capture) and a ServerStreamingProvider with a VLM ServerInputAdapter.

• Data (data/vlm_preprocessing.py, common/vlm_input.py): ShareGPT-style JSONL with an optional image / image_path / images field (path or base64; images takes the single list element). The target's own chat template renders the image into the first user turn; the target's HF processor expands the image region in id space, yielding expanded input_ids/loss_mask (trainer + passthrough) and collapsed request_input_ids + data:-URI image_data (capture request). One image per sample (v1); text-only samples can mix in the same run. Unreadable images and multi-image samples fail loudly (ImageDataError), never a silent text-only downgrade.

• Capture (inference/adapters/server_capture.py, training/disaggregated.py, launch_plan.py, patches/sglang/v0.5.14/spec-capture.patch): the patched server expands the image, runs the ViT, and writes aux hidden states straight to Mooncake; the capture request rides the stock input_ids + image_data /generate path. Managed capture servers set SGLANG_MM_AVOID_RETOKENIZE=1 for input_modality=multimodal so the server re-expands placeholders in id space with zero retokenization drift; client/server expansion mismatch fails loudly at the seq-len check. On Ascend hosts the capture servers default to --mm-attention-backend ascend_attn for non-text modalities (sdpa vision backend materializes [heads, N, N] scores and OOMs on large images); model.sglang_mm_attention_backend overrides.

• Draft model (modeling/draft/dflash.py, partial rotation ported from [Feature] VLM DFlash Training: Multi-Model Support for Qwen3-VL / Qwen3.5 / Qwen3.6 #585 by @zyk42, authorship preserved): partial rotation in apply_rotary_pos_emb (rotary_dim < head_dim, for Qwen3.5/3.6 partial_rotary_factor=0.25). The draft always uses the stock Qwen3RotaryEmbedding; [Feature] VLM DFlash Training: Multi-Model Support for Qwen3-VL / Qwen3.5 / Qwen3.6 #585's interleaved-mRoPE rotary is deliberately dropped (see Design).

• Target handling (modeling/target/target_utils.py, algorithms/model_providers.py): auto-select model.language_model.embed_tokens.weight for VLM targets; language-model depth via the text_config fallback for draft config auto-generation.

• Training (common/dflash_family_model.py): OnlineDFlashModel builds positions with the unconditional text-path 1D convention - multimodal batches take the byte-identical forward as text. Text-only runs are unchanged.

• Recipe: examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml, using the stock plain-rope configs/qwen3.5-4b-dflash.json. Warm-starting from the official draft via model.draft_checkpoint_path (weights-only) is supported out of the box.

• Tests/docs: tests/test_algorithms/test_dflash_multimodal.py (registration/provider gates, expansion math, request construction, 1D-positions regression); text-only assertions in test_config suites updated; status doc docs/advanced_features/vlm_dflash.md.

Scope notes (deliberate): DFlash only; one image per sample; offline (precomputed) multimodal capture not included; online evaluation for multimodal runs not included.

Related Issues

• Addresses [Feature] Support image-text multimodal input for DFlash training, similar to Eagle3 #583
• Related: [Feature] [RFC] DFlash Training Adaptation for Qwen VL Models #461 (DFlash × Qwen-VL RFC), Development Roadmap 2026 Q1 #374 (2026 roadmap - VLM support)
• Builds on [Feature] VLM DFlash Training: Multi-Model Support for Qwen3-VL / Qwen3.5 / Qwen3.6 #585 (draft model + configs, author preserved via cherry-pick)

Accuracy Test

• End-to-end (Ascend NPU, live run - PASSED): Qwen3.5-4B target + plain-rope DFlash draft, disaggregated managed run (5 capture servers + 5 trainer ranks, 32k context). Training runs end-to-end with healthy numerics: draft loss decays from ~5.7 to ~1, draft token accuracy rises from ~0.10 to ~0.65+, grad_norm stays bounded (0.4–1.1); accuracy_denom oscillates with image/text batch composition as expected for mixed VLM batches.

• Unit/integration (CPU): registration parity and provider gates, request/payload construction, image-expansion math, loud image errors, data-URI transport, golden topology/recipe tests - tests/test_algorithms/test_dflash_multimodal.py plus updated test_config suites pass.

• Patch integrity: the regenerated spec-capture.patch applies cleanly both ways to pristine sglang v0.5.14 (git apply --check / --reverse --check).

Benchmark & Profiling

Serving-side validation of the trained plain-rope draft (stock serving stack, no engine patches). Dataset: CC-OCR (~6.7K image-text samples; responses greedy-regenerated with the target; test split N=200, disjoint from train; image modality, K=7). Baseline = the official z-lab/Qwen3.5-4B-DFlash production draft served as-is.

Metric definitions (accept metrics use the server-side convention; the client-side per-request accept length is not reliably reported):

Metric	Source	Formula
Server accept rate (%)	server SpecDecoding metrics	Σ Accepted / Σ Drafted × 100
Server accept length (tokens)	server SpecDecoding metrics	Σ(mean_accept_len × Accepted) / Σ Accepted
Per-request throughput (tok/s)	client	Σ completion_tokens / Σ latency_s
Mean request latency (s)	client	Σ latency_s / n

Experiment 1: from-scratch (random init)

Draft trained from random initialization on the ~6.7K-sample train split, evaluated at step 2400 with 16 concurrent decoding instances.

Metric	From-scratch @ train	Untrained @ train	From-scratch @ test	Untrained @ test
Server accept rate (%)	38.25	34.70	28.39	34.96
Server accept length (tokens)	4.03	3.58	3.22	3.58
Per-request throughput (tok/s)	84.02	81.83	70.95	80.83
Mean request latency (s)	5.10	5.13	6.21	5.42

Findings:

• The plain-rope pipeline works end-to-end: train -> export -> serve, with the draft accepted by the stock serving path.
• Train-set accept slightly above baseline (+3.55pp) but with high variance; test-set clearly below baseline (-6.57pp accept, throughput -12%). Train-set up while test-set down is the classic memorization signature - at 6.7K samples we are squarely in [Feature] VLM DFlash Training: Multi-Model Support for Qwen3-VL / Qwen3.5 / Qwen3.6 #585's "<10K overfits severely" zone. The binding constraint is data scale and initialization, not the plain-rope route.

Experiment 2: warm-start (official draft init + in-domain fine-tune)

Draft warm-started from the official z-lab/Qwen3.5-4B-DFlash weights (model.draft_checkpoint_path, weights-only), fine-tuned on the same ~6.7K data, early-stopped at step 1000. Evaluated on the test split with 8 decoding instances per side. (Official-baseline numbers differ slightly from Experiment 1 due to different instance counts/batches; each table's comparison is same-batch.)

Metric	Warm-start @ test	Official untrained @ test	Δ
Server accept rate (%)	46.20	35.02	+11.18pp (1.32×)
Server accept length (tokens)	4.39	3.60	1.22×
Per-request throughput (tok/s)	98.77	82.49	1.20×
Mean request latency (s)	4.44	5.30	0.84×

Findings:

• Warm-start eliminates the from-scratch collapse: on the same test split, accept rate goes from 28.39% (from-scratch) to 46.20% - past the untrained production draft by +11.18pp.
• Official-weights warm-start + small in-domain fine-tune beats the production draft - large-scale data collection is no longer a prerequisite for a usable VLM draft; this is the recommended route at the current data scale.
• The trained plain-rope draft serves on the stock serving path with no engine patches - artifacts return directly to the official deployment surface.

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]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[curnane-lab]

Hi @jiapingW, could you take a review pass on this PR when free?