[RFC] DFlash

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

⚠️ 18 Awaiting Approval

As of commit 7261f77 with merge base d9e7a2e ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux Wheels (gh)
• Build documentation (gh)
• Build Linux Wheels (gh)
• Build macOS Wheels (gh)
• Build Presets (gh)
• Build Windows Wheels (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• MLX (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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.

[cthotti]

Hi @metascroy, thanks for the pointers.

On the C++ runner, I don't currently have access to hardware that could run it. I've spent the last few days trying to get an M4-M or M4-XL tier (the memory sizes gemma4-31b would need) through Scaleway, but they're either out of stock most of the time or too expensive. If anyone has spare M4-class hardware or access to something in that range, I'd genuinely like to help build this out further.

In the meantime, my hardware-side conclusions are written up in DLFASH_EXPERIMENTS.md in this PR, from which we learned that DFlash outperforms with up-to-date GPU architecture and on specific Math/Code prompts (not simple chat).

Also worth noting, the PR doesn't need everything I've touched locally to land. The parts that are actually useful here are:

• backends/mlx:
• dflash_draft_model.py
• export_llm_hf.py
• examples/models/qwen3:
• run_dflash.py
• run_baseline.py
• export_dflash_draft.py
• mlx_source_transformations.py
• DFLASH_EXPERIMENTS.md

For the next steps, I'm happy to open a separate issue for the gemma4-31b work specifically, if someone else would like to take a look.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 4m 18s —— View job

---

Code Review — DFlash on Qwen3-4B (updated)

Thanks for the follow-ups since my last review. I re-read the branch against origin/main. Good news first: the earlier blocking items are resolved — the dead Makefile target is gone (replaced with an explanatory comment), the tests were renamed test_dflash_* → check_dflash_* so pytest won't collect them, the profiling scaffolding is now behind --verbose, BSD headers are present, the typos are fixed, and the README now documents the manual/hardware-gated nature of the checks. The DFlash algorithm itself still reads as correct (draft→verify→first-mismatch→rollback, with the bonus-token in-bounds invariant now commented at run_dflash.py:169-172).

But re-reviewing surfaced one new blocking issue that dwarfs everything else.

🔴 Blocking — ~40,000 lines of generated MLX artifacts got committed by accident

Commit 9f86c90 ("Implementing Dflash_experiments.md for next users") added 159 files / 40,454 lines, of which only DFLASH_EXPERIMENTS.md (81 lines) and SUMMARY.txt (14 lines) are intentional. The other ~40,300 lines are auto-generated FlatBuffers code that is not tracked on origin/main and is meant to be produced at build time by backends/mlx/serialization/generate.py (see that dir's README.md — every one of these is marked [GENERATED]):

• backends/mlx/runtime/MLXLoader.cpp (2,473) + MLXLoader.h (2,743) + schema_generated.h (14,154)
• backends/mlx/serialization/_generated/ — 151 per-op node files + __init__.py
• backends/mlx/serialization/mlx_graph_schema.py (1,384) + _generated_serializers.py
• backends/mlx/_generated_inspector.py (966)

I confirmed none of these exist in origin/main (git cat-file -e origin/main:backends/mlx/runtime/MLXLoader.cpp → not found), while their sources (schema.fbs, generate.py, README.md) do. So they're regenerated locally and should never be checked in. This is what inflates the PR to "41,399 additions / 100+ files"; the actual DFlash change is ~1,000 lines across ~13 files.

Please git rm all of the generated files from the branch and confirm they're covered by the appropriate ignore rule so a local generate.py run won't re-stage them. That will also unblock a meaningful diff for reviewers. Remove generated artifacts →

🟡 Should fix

• SUMMARY.txt is stale and probably shouldn't be committed (examples/models/qwen3/SUMMARY.txt): it's an unstructured scratch dump of the M2-Air numbers, and it now contradicts DFLASH_EXPERIMENTS.md (0.57–0.84× vs the later 1.6–1.7× M4 results). Either fold the useful bits into DFLASH_EXPERIMENTS.md and delete SUMMARY.txt, or clearly label which machine it's from. Two overlapping perf files in one dir will confuse the next reader.
• DFLASH_EXPERIMENTS.md header date is inconsistent with the repo — it says Date: 08/16/2026, which is in the future relative to the PR. Minor, but worth correcting since it's a doc meant "for the next users."

🟢 Minor / nits (carried over — still open)

• export_dflash_draft.py:60 --ctx-len default of 8 is only the export-time example shape; fine given the dynamic ctx_dim, but a one-line note that this is just the tracing seed (not a runtime cap) would help, since --max-ctx-len is the real bound.
• run_dflash.py:189-190 EOS truncation: new_tokens[: new_tokens.index(eos_id) + 1] keeps the first EOS. Correct, but note that accepted_total/emitted_total are incremented before this truncation (:185-186), so the final-round τ/acceptance stats slightly over-count on the round that hits EOS. Harmless for benchmarking, but if you quote τ precisely you may want to count post-truncation.
• default_dflash_layer_ids unused — I flagged this last time; it looks like mlx_source_transformations.py no longer defines it (layer ids come from --dflash-layers), so consider this resolved unless it's still lurking elsewhere.

On the performance question (@metascroy)

Your DFLASH_EXPERIMENTS.md answers the earlier "no speedup?" thread well: base M4 shows 1.62× (math) / 1.70× (code), while M2/M2-Pro stay below 1.0×, and the write-up correctly attributes this to the Apple9 (M3/M4) GPU's matmul/dynamic-caching improvements making the block-verification pass cheaper — not to core count or bandwidth (the M2 Pro has more of both and still loses). The chat regression (τ≈2.9–4.2) is a draft-model quality limitation, independent of hardware, and is honestly documented. This is a solid, well-scoped RFC result. The MLX v32.0 pin bump (#20970) Scott mentioned would be worth re-benchmarking on top of, given it targets exactly the small-M gemm path that dominates verification here.

Summary

The DFlash implementation and its findings are in good shape and the prior review comments are addressed. The one thing blocking a clean merge is the ~40k lines of accidentally-committed generated FlatBuffers artifacts in 9f86c90 — remove those and the PR becomes the ~1k-line, reviewable change it's meant to be. SUMMARY.txt vs DFLASH_EXPERIMENTS.md duplication is a should-fix on top.
· branch dflash-qwen3-4b

[metascroy]

Is it possible to use the mlx rope method for this?

[cthotti]

I tried this, but ran into an issue as the existing rope op wants a fixed position offset, and ours needs to shift depending on context length, which broke the export. If you know a clean way to handle a moving offset like this, I'd be happy to implement it.

[metascroy]

I think F.rms_norm will lower more efficiently?

[cthotti]

Yea, nice idea. I swapped it over to F.rms_norm, kept the same order of operations so the math lines up exactly. And ran the lossless check and it's all good.

[cthotti]

DFlash on Gemma4-31B

Hi @metascroy,

I was able to get access to an M4-XL rental for a couple of days and used that time to test DFlash on Gemma4-31B. The results are very similar to what we saw with Qwen3: roughly 1.4-1.7x speedups on math and code prompts, but a slowdown on an open-ended chat (0.77x).

I also wanted to summarize what was implemented relative to the initial issue and design documentation you provided. The overall DFlash pipeline matches the intended design as well as the Qwen3 DFlash pipeline. But there were a few places where the implementation intentionally diverged from the design:

• The sample code in Part 1 still returns only the last-token logits, which works for normal autoregressive decoding but breaks DFlash verification. DFlash needs logits for every position in the drafted block, so the target export was modified to return full-sequence logits and hidden states
• Rather than creating a Gemma4-specific draft model, the existing generic DFlashDraftModel worked without modification because the released Gemma4 draft checkpoint follows the same architecture.
• The released draft checkpoint uses 6 tapped target layers [1, 12, 23, 35, 46, 57], so these were used
• Instead of integrating with LLMEngine, I wrote a small standalone C++ runner using the Executorch Module API. The goal was simply to determine whether Python/pybind overhead affected the results. Which it didn't, as the C++ and Python runners produced nearly identical performance across all prompt categories.

I've documented all the experiments in the dflash_experiments.md file.

For review, these are the most important files:

• dflash_export.py --> dflash target model and hidden-state export logic.
• export_dflash_target.py, export_dflash_draft.py--> export pipeline.
• run_dflash.py --> Python speculative decoding implementation.
• dflash_cpp_driver.cpp, CMakeLists.txt --> standalone C++ speculative decoding implementation (validation/benchmark tool, not full Gemma4_31BSpecDecEngine)

For future implementations, I'd be happy to work on properly integrating DFlash loop into Gemma4_31BEngine/LLMSession instead of leaving it as a standalone driver. Also happy to dig into anything else if you'd like.

[cthotti]

Hi, I accidentally closed this issue by renaming the branch, but its restored now.

[metascroy]

Thanks @cthotti! Let me find some time to check out your PR tomorrow or Friday and run some benchmarks myself.

[metascroy]

@cthotti I think this needs another pass before landing. The basic approach works, but the PR currently mixes a Qwen bringup, a second Gemma implementation, a generic export API, and a new draft transformer implementation. That makes the change much larger than necessary and leaves important correctness and performance gaps.

My recommendation is to narrow this PR to one generic DFlash path built on export_llm_hf.py, with Qwen as the first supported checkpoint.

Requested changes

1. Make block_len dynamic

The exported draft currently fixes the proposal token dimension to block_size. The runner cannot shorten the last round, so it can emit past max_new_tokens or run past the remaining context capacity.

Please export the proposal dimension as a bounded dynamic dimension and choose the runtime length from the remaining token and context budgets:

block_len = min(
    checkpoint_block_size,
    max_new_tokens - num_generated,
    max_seq_len - position,
)

The context and proposal dimensions should remain symbolically related throughout position construction and attention. Prefer deriving positions inside the draft wrapper so callers do not have to provide a third tensor with a coupled shape. If block_len=1 does not lower cleanly, use a target only step when one token remains.

I verified that the reference HF draft exports with dynamic context and proposal lengths. MLX lowering also works after adding a small operator.neg handler that maps to the existing integer multiplication by -1.

2. Reuse the existing transformer implementation

DFlashDraftModel reimplements RoPE, normalization, attention, masking, and the transformer blocks. This has already created behavior differences.

First, embed_scale exists but is never loaded. That is harmless for Qwen but wrong for model families that scale embeddings.

Second, the sliding window path adds a causal proposal mask even though the Qwen checkpoint targeted by this PR uses noncausal full attention.

Third, the custom attention calls PyTorch SDPA directly instead of using the Hugging Face attention dispatch that the MLX integration already extends.

Please adapt the checkpoint's HF reference module instead. It already reuses the HF Qwen blocks and attention registry, so behavior remains with the model implementation rather than being duplicated in ExecuTorch. A thin wrapper can still copy the target embedding and language model head weights to make the exported draft self contained.

Please also add an eager parity test against the reference module for a small input. This will catch future differences in masks, positions, and logits.

3. Include a real draft KV cache

The current draft reprojects the full accumulated target context on every round. Draft cost therefore grows with context length and becomes quadratic over a generation, while the target remains incremental. This undermines the performance case for speculative decoding on longer inputs.

Please export the draft with persistent cache state rather than treating cache support as a later optimization. Creating a new DynamicCache for every call does not solve this because it is removed while tracing the exported graph.

I verified that a small draft specific static cache wrapper can be exported and lowered through MLX. It follows the same pattern as TorchExportableModuleWithStaticCache: register the cache tensors as mutable buffers and pass cache_position through the draft call. The remaining work is to mask unused cache slots correctly and implement the DFlash append and crop position semantics after verification.

4. Narrow and consolidate the PR

Please scope this change to the export_llm_hf.py path and remove the parallel Gemma implementation from this PR. The Gemma files add another target exporter, another draft exporter, two more speculative loops, and artifact settings duplicated in command line flags.

The Qwen draft exporter and runner are already mostly generic. Their Qwen behavior is limited to defaults and prompt or tokenizer handling. Move the generic DFlash pieces next to the generic MLX LLM exporter and keep model behavior in a small adapter.

The intended structure should contain one generic target export path with optional hidden state taps, one generic draft export path, one speculative decoding loop, and thin checkpoint and tokenizer adapters.

This avoids establishing separate Qwen, Gemma, Python, and C++ DFlash implementations that will drift.

5. Make hidden state tapping an explicit generic feature

The generic exporter currently exposes this as --dflash-layers, but the mechanism simply returns selected hidden layers. Please rename it to something based on the mechanism, such as --tap-layers or --output-hidden-layers.

A sliding window model currently selects the hybrid cache wrapper before the tap layer option is considered. This silently exports a model without the requested hidden outputs. Since composing hidden tapping with the hybrid wrapper is outside the proposed scope of this PR, reject that combination explicitly instead of ignoring the option.

6. Fix the validation scripts

The target smoke check passes three tokens with one cache position even though those dimensions share one exported symbol. It should pass one position per token and exercise more than one sequence length.

The losslessness script ignores subprocess failures and exits successfully when outputs differ. Two failed runs can therefore compare as empty output and report PASS.

Please make command failures and output mismatches fail the check. Also add small tests for the acceptance loop using synthetic inputs. The cases should include rejection of the first token, full acceptance with a bonus token, EOS handling, and a remaining generation budget smaller than the native block.

Cleanup

Please rename or remove experiment files that are not part of the supported workflow and clean the trailing whitespace in the added Python and Markdown files.

Landing criteria

I would be comfortable landing once the PR has dynamic proposal length with generation and context bounds, the HF reference draft module, persistent draft KV caching, one generic export and runner path, explicit hidden tapping behavior, and reliable validation scripts. Gemma support can follow after the generic path is established.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 29s —— View job

---

Code Review — DFlash (Qwen3-4B + Gemma4-31B)

I re-read the branch at 1f0894b against the base. Since the last review, the two most recent commits are lint-only (41de477, 1f0894b), so this pass checks the current state against your 2026-08-03 landing criteria and flags what's still open. The core speculative loop still reads as correct (draft → verify → first-mismatch → rollback, bonus-token invariant intact at run_dflash.py:169-172), but the six requested changes are essentially all still outstanding.

Landing criteria — status

1. Dynamic block_len — ❌ not done. The proposal token dim is still hard-fixed. export_dflash_draft.py:92 exports "tokens": None (static), and run_dflash.py:139 hard-codes bs = block_size with a comment explicitly acknowledging it can't shorten the last round. So the runner can still emit past max_new_tokens / run past context. The block_len = min(checkpoint_block_size, max_new_tokens - num_generated, max_seq_len - position) bound you asked for isn't wired in. Fix →

2. Reuse the HF reference module — ❌ not done. dflash_draft_model.py still reimplements RoPE, RMSNorm, attention, masking, and the decoder block (277 LOC). The three concrete divergences you called out are all still present:

• embed_scale is defined on DFlashConfig (:59) and applied at :230, but load_dflash_config (:257-277) never populates it — so it's always 1.0. Harmless for Qwen, wrong for Gemma-style embedding scaling.
• _sliding_mask (:184-190) applies a causal proposal mask (k_pos <= q_pos), whereas the Qwen checkpoint targeted here uses full non-causal proposal attention.
• Attention calls torch.nn.functional.scaled_dot_product_attention directly (:179) instead of routing through the HF attention dispatch the MLX integration extends.

No eager-parity test against the reference module exists yet.

3. Real draft KV cache — ❌ not done. run_dflash.py:203 still does hidden = torch.cat([hidden, new_hidden...]) and the draft re-projects the full accumulated context every round (dflash_draft_model.py:232 self.fc(target_hidden) over the whole target_hidden). Draft cost stays O(context) per round → quadratic over a generation while the target is incremental. No persistent draft cache buffers were added.

4. Narrow / consolidate — ❌ not done. The parallel Gemma implementation is still fully present: examples/models/gemma4_31b/{dflash_export.py, export_dflash_target.py, export_dflash_draft.py, run_dflash.py, run_baseline.py, dflash_cpp_driver.cpp} plus CMakeLists.txt. That's a second target exporter, second draft exporter, second speculative loop (Python), and a C++ loop — the exact duplication the request asks to remove until the generic path lands. Net diff is still ~2,168 lines across both models rather than one generic Qwen-first path.

5. Explicit hidden-tapping feature — ❌ not done. Still exposed as --dflash-layers (export_llm_hf.py:458), not renamed to a mechanism-based name. And the ordering hazard you flagged is real: in _export_with_custom_components, the branch is if sliding_window … elif dflash_layers … (:212 vs :223), so a sliding-window model silently ignores --dflash-layers and exports without hidden outputs instead of rejecting the combination. Fix →

6. Validation scripts — ❌ not done. Both problems you named are still there:

• check_dflash_lossless.py:18-33 reads only .stdout and never checks the subprocess return code; :42 compares baseline.strip() == dflash.strip(). Two failed runs both produce empty output → compare equal → prints PASS. Fix →
• check_dflash_target.py:28-29 passes 3 tokens with a 1-element input_pos (torch.tensor([0])), even though the target export ties input_ids seq-len and cache_position to one symbol (export_llm_hf.py:285-289). Should pass one position per token and exercise more than one sequence length.
• No synthetic-input tests for the acceptance loop (first-token rejection, full-accept + bonus, EOS handling, budget < native block) were added.

Cleanup — ❌ not done

• Trailing whitespace remains in the added files (13 occurrences under examples/models/qwen3/, 11 under examples/models/gemma4_31b/). lintrunner didn't strip these because these dirs likely aren't in the lint config's include set — worth a manual pass.
• Experiment scaffolding (DFLASH_EXPERIMENTS.md in both dirs, dflash_cpp_driver.cpp as a standalone benchmark tool) is still in the tree.

One correctness nit worth a comment

run_dflash.py:189-191: when EOS lands, accepted = min(accepted, len(new_tokens) - 1) correctly clamps, but note the draft-context growth at :203 uses new_hidden[:, :len(new_tokens), :] after truncation — that's consistent and correct, just subtle. Fine to leave, but a one-liner would help since an off-by-one here silently corrupts the next round's conditioning.

Summary

The algorithm and findings are solid, and the write-ups answer the earlier perf question well (base M4 1.4–1.7× on math/code). But against the six landing criteria, all six are still open — the branch head only carries lint fixes since the request. The highest-leverage items to unblock a clean, reviewable merge are: (2) adapt the HF reference draft instead of the hand-rolled 277-LOC module, (3) add the persistent draft KV cache, and (4) drop the Gemma duplication so this is one generic Qwen-first path. (1), (5), and (6) are smaller and mostly mechanical.
· branch dflash-qwen3-4b