[RFC] DFlash on Qwen3-4b

[cthotti]

Summary

This PR adds support for DFlash speculative decoding to the ExecuTorch MLX delegate using Qwen3-4B as the reference model. It has hidden-state export for the target model, a configurable DFlash draft model, export support for both models, and an end-to-end speculative decoding driver. The implementation has been verified to produce an identical output to standard greedy decoding.

Although the original issue was targeted at Gemma, my implementation was developed and validated on Qwen3 because publicly available DFlash draft weights are available and the model fits within local hardware constraints. The overall implementation is model-agnostic and designed to support additional architectures in follow-up work.

Fixes #20701

DFlash Overview

Unlike conventional speculative decoding, where a small autoregressive draft model predicts one token at a time, DFlash predicts an entire block of tokens in a single forward pass over a masked block. The draft model is conditioned on intermediate hidden states from the target model, which are added into the Key/Value projections of every draft rather than only at the input. This provides richer conditioning throughout the draft network while keeping the draft model relatively small.

The draft model is implemented through a configurable DFlashConfig, making it adaptable to different transformer architectures. It has been numerically verified on Qwen3 and structurally validated against Llama-3.1 and Gemma-style configurations. There are additional Gemma-specific features documented, like partial RoPE and post-layer scaling, but they remain outside the scope of this PR.

Key Implementation Details

Phase 1: Target hidden-state export

• Adds TorchExportableModuleWithStaticCacheAndHidden, which extends the existing Hugging Face export wrapper to return both logits and concatenated hidden states from the selected target layers.

Phase 2: Draft model and export

• Adds a configurable DFlash draft model together with its export pipeline.
• The exporter supports dynamic context lengths and 4-bit quantization compatible with the existing MLX export flow.

Phase 3: Speculative decoding driver

• Adds an end-to-end Python implementation of the DFlash decoding algorithm (run_dflash.py), including:
• Draft block generation, Target verification, First-mismatch acceptance, and KV-cache position rollback
• Applies the Qwen3 chat template with enable_thinking=False, matching the format used during draft model training.

Python Driver

My PR includes the Python implementation of the speculative decoding loop rather than the earlier C++ engine recommended in the issue.

Profiling showed that nearly all execution time is spent inside the exported target and draft model execution, while the Python overhead contributes only a very small fraction of the total runtime. Since both implementations execute the same exported MLX programs, the Python driver provides similar functionality and is simpler to maintain. A C++ implementation can be added later if tighter runtime integration is needed.

Performance

My primary goal of this PR is to add DFlash support and verify that the implementation is correct. In addition to correctness, I benchmarked the implementation on Apple M2 hardware to understand its performance characteristics.

The most significant observation is that the target model's verification pass scales almost linearly with the number of tokens being verified. On this system, verifying a single token (T=1) takes approximately 32 ms, while verifying a block of 16 tokens (T=16) takes approximately 239 ms, roughly 7.5× the latency for a 16× larger verification window.

Since speculative decoding relies on verifying multiple draft tokens in a single target forward pass, this scaling behavior directly impacts the achievable speedup. The measured acceptance rates (chat: 4.17, math: 5.59, code: 6.81) are not sufficient to offset the additional verification cost on this hardware, so the implementation is currently slower than baseline greedy decoding for a single-request inference on an Apple M2.

This appears to be because of the underlying hardware rather than the DFlash implementation itself. The speculative decoding algorithm is functioning correctly and produces token-for-token identical output to baseline decoding. Platforms where verification cost grows more slowly with batch size are expected to benefit more from DFlash, since speculative decoding becomes increasingly effective as the cost of verifying larger blocks approaches the cost of verifying a single token.

Test Plan

All tests were run using exported target and draft models (qwen3_4b_dflash_target.pte and qwen3_4b_dflash_draft.pte) with a block size of 16 and 4-bit quantization.

Lossless speculative decoding:
python3 examples/models/qwen3/tests/test_dflash_lossless.py
Verifies that DFlash produces token-for-token identical output to standard greedy decoding.

Target model export

python3 examples/models/qwen3/tests/test_dflash_target.py \
    qwen3_4b_dflash_target.pte

Draft model export and dynamic context lengths

python3 examples/models/qwen3/tests/test_dflash_draft.py \
    qwen3_4b_dflash_draft.pte

Final Results

Math prompt

• baseline_runs (tok/s) = [25.73, 24.63, 25.65]
• baseline_median = 25.65
• dflash_runs (tok/s) = [19.44, 19.47, 19.44]
• dflash_median = 19.44
• τ = 5.59
• speedup = 0.76×

Code prompt

• baseline_runs (tok/s) = [28.28, 28.19, 26.10]
• baseline_median = 28.19
• dflash_runs (tok/s) = [23.84, 23.81, 22.70]
• dflash_median = 23.81
• τ = 6.81
• speedup = 0.84×

Chat prompt

• baseline_runs (tok/s) = [26.05, 26.11, 27.98]
• baseline_median = 26.11
• dflash_runs (tok/s) = [14.78, 14.77, 14.75]
• dflash_median = 14.77
• τ = 4.17
• speedup = 0.57×

cc @mergennachin @iseeyuan @lucylq @helunwencser @tarun292 @kimishpatel @jackzhxng

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20840

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-cla]

Hi @cthotti!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: cthotti / name: cthotti (04ae566, 37f764c, 5b19ffc, d2b9970)