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

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[metascroy]

@claude review this code

[metascroy]

Thanks @cthotti! Am I right in reading the PR that you're not seeing speedup from dflash?

[claude]

Claude finished @metascroy's task in 3m 33s —— View job

---

Code Review — DFlash on Qwen3-4B

Nice work — the DFlash design is well thought out, the model-agnostic DFlashConfig is clean, and the speculative loop is correct (I traced the draft/verify/bonus-token indexing; it does produce the lossless bonus token). Below are the issues I found, grouped by severity.

🔴 Blocking / correctness

1. Makefile target references gitignored, non-existent files
Makefile adds a qwen3_dflash-mlx target that runs cd examples/models/qwen3 && cmake --workflow --preset qwen3-dflash-mlx. But examples/models/qwen3/CMakePresets.json, CMakeLists.txt, and the C++ engine sources are all .gitignored in this same PR (commit 7be25cd, marked "incomplete, broken build, follow-up PR"). So this target can never succeed for anyone who checks out the branch — the preset file it needs is intentionally excluded from the repo. Please drop the Makefile target (and the corresponding .gitignore C++ block) until the follow-up PR that actually lands the engine. Shipping a dead build target is confusing. Fix this →

2. tests/test_dflash_*.py will crash if pytest ever collects them
test_dflash_target.py:17 (pte_path = sys.argv[1]) and the others execute logic at module top level, not inside test functions. If pytest imports these (e.g. pytest -n auto examples/models/qwen3, or if someone adds the dir to testpaths), collection fails immediately: test_dflash_target.py raises IndexError on sys.argv[1], and all three try to load .pte files / download from HF / require Apple MLX hardware at import time. They're really manual driver scripts, not pytest tests. Either (a) rename them to run_*/check_* so they aren't matched by pytest's test_* glob, or (b) wrap the bodies in real def test_...() functions guarded with pytest.importorskip/skip-if-no-pte so they degrade gracefully. As-is they can only be run by hand with args, which the naming hides.

🟡 Should fix

3. Leftover profiling scaffolding in run_dflash.py
The main loop (run_dflash.py:144-201) is heavily instrumented with _t0, _draft_exec_time, _draft_argmax_time, _target_argmax_time, _fm_time, _cat_time, etc., interleaved with the real logic. This obscures the algorithm and conflicts with the repo's "minimal comments / self-documenting" guidance. Since the argmax/first-mismatch/cat timings are all sub-millisecond (as the PR body notes, time is dominated by model execution), consider removing most of this and keeping only the draft_exec/target_exec timings behind --verbose.

4. Missing BSD license headers
dflash_draft_model.py, export_dflash_draft.py, mlx_source_transformations.py, run_dflash.py, run_baseline.py, and the three test files all start with a module docstring but no Copyright (c) Meta Platforms BSD header that the rest of the repo uses (see export_llm_hf.py:1-6). The LICENSELINT rule only rejects proprietary headers rather than requiring the BSD one, so lint won't catch it, but please add them for consistency.

🟢 Minor / nits

• --chat-template is a no-op flag (run_dflash.py:46-51, run_baseline.py:18): action="store_true", default=True means passing --chat-template changes nothing — only --no-chat-template has effect. Drop the store_true half and keep just the --no-chat-template toggle.
• Typos: frmo → from (dflash_draft_model.py:151); prodcues → produces (test_dflash_lossless.py:1); an .pte → a .pte and output project weights → output projection weights (export_dflash_draft.py:1,3).
• first_mismatch is correct but subtle (run_dflash.py:30-35): it relies on len(draft_ids) == block_size-1 while target_ids has block_size entries, so target_ids[accepted] is always in-bounds (including the all-accepted bonus-token case). A one-line comment noting the length asymmetry would help future readers, since an off-by-one here would silently break correctness.
• default_dflash_layer_ids (mlx_source_transformations.py:66-70) is unused in the committed code path — the layer ids come from the --dflash-layers CLI arg / the draft config. If nothing calls it, consider removing it or wiring it up as the CLI default.

Testing / CI note

None of the DFlash tests can run in this repo's CI: they require exported qwen3_4b_dflash_target.pte / _draft.pte (multi-GB, not checked in), HF downloads, and Apple M-series MLX hardware. That's understandable for an MLX-only feature, but it means the "lossless" guarantee is only manually verified. Worth stating explicitly in examples/models/qwen3/README.md (which I didn't see updated) so future contributors know these are manual, hardware-gated checks.

---

Overall the core algorithm and export pipeline look solid and the correctness claim checks out on inspection. The main things to resolve before merge are the dead Makefile/gitignore C++ references (#1) and the pytest-collection hazard (#2).
· branch dflash-qwen3-4b

[cthotti]

Thanks @cthotti! Am I right in reading the PR that you're not seeing speedup from dflash?

Yes, I didn't see a speedup on my original hardware, although the implementation is functionally correct and produces token for token identical output to baseline decoding. From profiling, the verification stage dominated each speculative round on my Macbook M2 Air (8 core GPU), so the cost of verifying each drafted block outweighed the time saved by drafting.

However, since then, I've rerun the implementation on another Apple Silicon environment using a MacBook M4 (10-core GPU), where I observed a 1.4x speedup over a baseline run. Which makes me think that earlier results were more so hardware dependent. Before drawing a conclusion, I'd like to spend more time running experiments and testing out other Apple environments, like a M2 Pro (16 core GPU) and, if possible, M4 Pro (20 core GPU).

Would it be alright if I spend a little more time validating the implementation and updating the PR with those results?

[metascroy]

Thanks @cthotti! Am I right in reading the PR that you're not seeing speedup from dflash?

Yes, I didn't see a speedup on my original hardware, although the implementation is functionally correct and produces token for token identical output to baseline decoding. From profiling, the verification stage dominated each speculative round on my Macbook M2 Air (8 core GPU), so the cost of verifying each drafted block outweighed the time saved by drafting.

However, since then, I've rerun the implementation on another Apple Silicon environment using a MacBook M4 (10-core GPU), where I observed a 1.4x speedup over a baseline run. Which makes me think that earlier results were more so hardware dependent. Before drawing a conclusion, I'd like to spend more time running experiments and testing out other Apple environments, like a M2 Pro (16 core GPU) and, if possible, M4 Pro (20 core GPU).

Would it be alright if I spend a little more time validating the implementation and updating the PR with those results?

Yes, that would be great! Pybindings do add overhead, and we do usually benchmark with c++ runners.

We have a C++ runner for gemma4-31b (this would require an M4 to run).

[metascroy]

@cthotti FYI, MLX had some upstream improves for small M gemm in v32.0. Here is an ET pin bump: #20970

You might see better perf with this on top of your PR.