Dflash: Block Diffusion Speculative Decoding by vjanfaza · Pull Request #995 · quic/efficient-transformers

vjanfaza · 2026-05-18T16:44:20Z

In this PR, we are adding the support for DFlash (Block Diffusion for Flash Speculative Decoding) approach that significantly accelerate LLM inference on Qualcomm platforms.
DFlash replaces autoregressive draft models with a diffusion-based block drafter that predicts an entire block of tokens in a single parallel pass, resulting in constant drafting latency independent of block size.

Key benefits:

5–6× lossless speedup over baseline decoding
~2.5× faster than EAGLE-3
Higher acceptance length enabled by KV injection of target hidden states into every draft layer
Improved accelerator utilization and reduced end-to-end latency

Verification is performed by the frozen target model, guaranteeing output identical to standard decoding.

Use-case(s)

High-throughput and low-latency LLM serving
Long chain-of-thought reasoning workloads (e.g., math and planning)
Code generation and agentic inference loops
Platform differentiation via advanced speculative decoding on AI100

#Requirements
DFlash speculative decoding with block diffusion drafting and target verification
Extraction and fusion of hidden states from multiple target layers during prefill
KV cache injection of fused target context into all draft layers
Integration with QEfficient, vLLM, and other applications
Support current architecture:
Target: single QPC for prefill + decode (expandable to DA serving)
Draft: lightweight QPC

Equally Contributors

quic-rishinr · 2026-05-25T09:03:40Z

Design discussion pending.

Add DFlash speculative decoding: a block-diffusion draft language model (DLM) that proposes block_size tokens per step from the target model's hidden states, verified by the target language model (TLM). QEfficient side: - Qwen3 DFlash draft model (two-stream attention over target_hidden + noise). - DFlashTransform swaps in the DLM modules; DFlashDLMTransform injects lm_head/embed_tokens from the paired TLM checkpoint; DFlashTLMTransform attaches fc/hidden_norm (weights from the DLM checkpoint, fc scaled for fp16 range) and activates target hidden-state collection. Callers select the paths via qaic_config (dflash_dlm / target_layer_ids) plus the cross-checkpoint repo keys (dflash_tlm_repo / dflash_dlm_repo); no model surgery on the caller side. - modeling_auto: DFlash export inputs/outputs, block_size specialization, and batched DLM decode compile. examples/performance/dflash: single-prompt and benchmark runners plus a make_models helper, driven through QEFFAutoModelForCausalLM.from_pretrained. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com> # Conflicts: # QEfficient/transformers/models/modeling_auto.py

DFlashDLMTransform/DFlashTLMTransform.apply() use torch.no_grad()/torch.load, but the bare import torch was dropped when merging upstream main (which no longer needed it elsewhere in the file), leaving only from torch import nn. Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

Extends DFlash speculative decoding to vision-language models. gemma4 and qwen3-vl TLMs collect target-layer hidden states, attach fc/hidden_norm, and emit them as an extra output for the DFlash draft (DLM) to verify against. DFlashTLMTransform generalizes to VLM configs (nested text_config/language_model) alongside the existing dense-model path. utils.py adds VLM-aware DLM weight injection (compile_gemma_vlm_dlm_qpc) since a VLM base keeps lm_head/embed_tokens on a nested submodule that the flat-key DFlashDLMTransform loader can't reach; extracting and injecting them explicitly fixes drafts never being accepted (AR stuck at 1.00) and restores the expected acceptance rate. New vision example scripts (basic_inference_vision.py, benchmark_vision.py, dflash_spd_vision_single_prompt.py, dflash_spd_vision_benchmark.py, dflash_spd_vision_text_benchmark.py) run the vision-encoder + DFlash-aware language QPC + draft QPC pipeline for single-prompt and dataset-benchmark use, covering text-only, image+text, and language-only-on-text-dataset modes. Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com> Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

quic-rishinr · 2026-08-25T19:14:20Z

Please add unit test, full end to end test and update the documentation as well

… draft model classes, refactor example scripts - Extract the DFlash TLM target-layer hidden-state computation (concat -> fc -> hidden_norm) shared by llama/qwen3/gemma4/qwen3_vl into QEfficient/transformers/spd/dflash.py, reused via compute_dflash_target_hidden_states. - Fix a dict-overwrite bug in the DFlash draft model's ONNX export (QEfficient/transformers/models/modeling_auto.py) that silently dropped input_ids/position_ids from dynamic_axes, baking a static seq_len into the exported graph and causing a runtime buffer-size mismatch. - Rename the DFlash draft model wrapper classes from QEffQwen3* to QEffDFlash* to avoid confusion with the real Qwen3 modeling classes. - Move the DFlash generation/SPD loops out of examples/performance/dflash/utils.py and the deleted single-prompt scripts into QEfficient/generation/dflash_generation.py, exposing run_text_inference/run_vision_inference as direct in-process entry points. - Rename basic_inference.py to basic_inference_text.py; basic_inference_text.py and basic_inference_vision.py now call dflash_generation directly instead of shelling out to dflash_spd_single_prompt.py/dflash_spd_vision_single_prompt.py (both removed). - Remove the now-dead DFlash target_hidden_states rebinding block from kv_offload_generate in modeling_auto.py (never exercised by any dflash example). - Update examples/performance/dflash/README.md to match the current script layout. Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com> Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

…che sizing - QEffDFlashRotaryEmbedding.forward sliced cos/sin to [:seq_len] before returning, but callers gather by absolute position and DFlash's noise positions run up to 2x kv_seq_len ahead of the context; return the full cached table instead of truncating it. - Size the dummy KV cache to seq_len*2 for dflash_dlm exports so it covers both the context and noise-position halves of the window. Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com> Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

Adds a CPU-only unit test suite for DFlashTransform/DFlashDLMTransform/ DFlashTLMTransform and compute_dflash_target_hidden_states, and a full end-to-end test on real QAIC hardware that drives run_spd_inference_single through a tiny random-weight TLM/DLM pair. Also documents DFlash in the supported features table. Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com> Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

quic-rishinr · 2026-09-02T11:38:57Z

@@ -0,0 +1,7 @@
+# DFlash example run artifacts (not source)


please remove this file

Confirmed that .gitignore matches upstream main and is no longer present in the current PR diff.

it is part of the current PR diff

Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com> Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

quic-rishinr · 2026-09-08T09:43:45Z

+    mm_token_type_ids: np.ndarray | None = None,
+    vision_embeds: np.ndarray | None = None,
+) -> SpecDecodingMetrics:
+    batch_size = 1


if we are supporing only batch_size=1 we should add a fallback warning if user is passing batch size > 1

We have the support of higher batch sizes and I tested it in vLLM with no issue. I fixed this part of the code

Then in that case we should be remove hardcoding batch_size to 1

Fixed by adding fallback as you suggested. We have it in vLLM so in the fallback warning we are suggesting to use vLLM for higher batch sizes since the current SpD implementation in QEfficient for higher batch sizes is not very efficient and it's just repeating it based on batch size which is not very effective way to do it. However, in vLLM we have an efficient implementation of it which the results can be found in this link

quic-rishinr · 2026-09-08T10:18:23Z

+            cache_kwargs = {"batch_index": batch_index, "position_ids": position_ids_target}
+            if comp_ctx_lengths is not None:
+                attention_mask = attention_mask[:, :, :, : comp_ctx_lengths.shape[-1]]
+                cache_kwargs["CCL"] = attention_mask.shape[-1]


Was the CCL tested in this case? since write_only() only reads position_ids and batch_index, it ignores CCL in Dynamic write_only() method, So it will fall back to full ctx length.

CCL doesn't have any effect on write_only and it only helps during gathering which happens in update and read_only functions

yes, but CCL is passed to write_only(), which ignores it and then cache_kwargs is updated in line 276 "cache_kwargs = {"batch_index": batch_index, "position_ids": position_ids_target}" before update(). This way update() never receives CCL and falls back to the full cache length.

Was this tested with CCL?

Very good point. We tested it with CCL before but later with some newer changes we missed this point. Thank you for pointing it out. Fixed it.

Introduce the `dflash2_draft` package (the DFlash-2 draft model, which has no upstream `transformers` implementation, so its classes are written export-ready and used directly) together with the wiring that selects it: the new `DFlash2DLMTransform` in `pytorch_transforms.py` and the `qaic_config['dflash2_dlm']` path in `modeling_auto.py`. The `dflash_dlm` and `dflash2_dlm` flags are mutually exclusive, and DFlash-2 shares DFlash's doubled KV-cache sizing, `target_hidden` example input, and block-size specialization. For the existing DFlash draft, move the rotary embedding from per-attention to the model level: `QEffDFlashModel.__qeff_init__` builds the table once and folds `attention_scaling` into the `sin_cached`/`cos_cached` Parameters, which are then threaded through the decoder layers into attention. This removes the per-layer rotary instance and its `forward`. The whole table is passed down because `qeff_apply_rope_two_streams` gathers it once per stream — `position_ids_target` for the context keys and `position_ids` for the noise keys — so a pre-gather at the model level would drop the target stream. Only the modeling files are in scope here; the compilation and TLM paths are not complete yet. Co-authored-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: FNU Annanya <fannanya@qti.qualcomm.com>

Signed-off-by: FNU Annanya <fannanya@qti.qualcomm.com>

quic-rishinr · 2026-09-11T06:10:27Z

@vjanfaza Let's not increase the scope of this PR. Adding DFlash2 changes here will only delay the merge further.

quic-rishinr · 2026-09-11T06:13:46Z

+    mm_token_type_ids: np.ndarray | None = None,
+    vision_embeds: np.ndarray | None = None,
+) -> SpecDecodingMetrics:
+    batch_size = 1


Then in that case we should be remove hardcoding batch_size to 1

quic-rishinr · 2026-09-11T06:33:14Z

+                cache_kwargs["CCL"] = attention_mask.shape[-1]
+
+            # first write for target positon_id
+            past_key_value.write_only(k_ctx, v_ctx, self.layer_idx, cache_kwargs)


was it benchmarked? Can you share the latency numbers?

quic-rishinr · 2026-09-11T06:40:14Z

+            cache_kwargs = {"batch_index": batch_index, "position_ids": position_ids_target}
+            if comp_ctx_lengths is not None:
+                attention_mask = attention_mask[:, :, :, : comp_ctx_lengths.shape[-1]]
+                cache_kwargs["CCL"] = attention_mask.shape[-1]


yes, but CCL is passed to write_only(), which ignores it and then cache_kwargs is updated in line 276 "cache_kwargs = {"batch_index": batch_index, "position_ids": position_ids_target}" before update(). This way update() never receives CCL and falls back to the full cache length.

Was this tested with CCL?

quic-rishinr · 2026-09-11T06:43:34Z

@@ -0,0 +1,7 @@
+# DFlash example run artifacts (not source)


it is part of the current PR diff

Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com> Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com>

vjanfaza force-pushed the DFlash branch from 35e0a5c to 3b04208 Compare May 18, 2026 16:47

quic-rishinr requested review from jinkasaisagar, quic-amitraj, quic-rishinr and vbaddi and removed request for jinkasaisagar and quic-amitraj May 19, 2026 04:15

anujgupt-github marked this pull request as draft May 19, 2026 10:31

vjanfaza force-pushed the DFlash branch from 0fb7610 to 68a9284 Compare May 22, 2026 20:41

vjanfaza force-pushed the DFlash branch 3 times, most recently from 630d915 to 4ec6a39 Compare May 29, 2026 04:59

vjanfaza force-pushed the DFlash branch 6 times, most recently from ce1f9b7 to db57abd Compare June 25, 2026 20:54

vjanfaza force-pushed the DFlash branch from db57abd to 8c98c32 Compare June 25, 2026 21:17

vjanfaza and others added 3 commits August 13, 2026 17:19

Merge branch 'main' into DFlash

c8beca8

Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com> # Conflicts: # QEfficient/transformers/models/modeling_auto.py

vjanfaza marked this pull request as ready for review August 24, 2026 18:11

quic-rishinr requested changes Aug 25, 2026

View reviewed changes

vjanfaza and others added 2 commits August 27, 2026 14:16

Merge remote-tracking branch 'upstream/main' into DFlash

9941254

Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

vjanfaza force-pushed the DFlash branch from 1893647 to 9941254 Compare August 27, 2026 21:16

vjanfaza and others added 2 commits August 27, 2026 15:24

quic-rishinr requested changes Sep 2, 2026

View reviewed changes

vjanfaza and others added 3 commits September 3, 2026 17:12

Refactor DFlash generation and examples

a12cf54

Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com> Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

Harden DFlash transforms and checkpoint loading

13459d3

Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com> Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

Merge upstream main into DFlash

a8e49e1

Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com> Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com>

vjanfaza force-pushed the DFlash branch 2 times, most recently from 1e4ef14 to a8e49e1 Compare September 4, 2026 00:59

quic-rishinr requested changes Sep 8, 2026

View reviewed changes

fannanya and others added 2 commits September 9, 2026 19:38

Removing unsupported models form utils

4bee43d

Signed-off-by: FNU Annanya <fannanya@qti.qualcomm.com>

quic-rishinr requested changes Sep 11, 2026

View reviewed changes

Remove DFlash2 support from DFlash PR

3fa3c73

Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com> Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com>

vjanfaza force-pushed the DFlash branch from 88e20f6 to 3fa3c73 Compare September 11, 2026 19:46

vjanfaza and others added 2 commits September 11, 2026 15:40

Address DFlash review comments

f5eb64f

Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com> Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com>

Add batch size handling to DFlash examples

c7ff0f3

Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com> Co-authored-by: FNU Annanya <fannanya@qti.qualcomm.com>

Uh oh!

Conversation

vjanfaza commented May 18, 2026

Key benefits:

Use-case(s)

Equally Contributors

Uh oh!

quic-rishinr commented May 25, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

quic-rishinr commented Aug 25, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vjanfaza Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

quic-rishinr commented Sep 11, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

vjanfaza Sep 11, 2026 •

edited

Loading