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

❌ 3 Pending, 1 Unclassified Failure

As of commit e3156c4 with merge base d827cf7 ():

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• MLX / test-mlx-stories110m / test-mlx-stories110m (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  ##[error]fatal: unable to access 'https://github.com/nlohmann/json.git/': Failed to connect to github.com port 443 after 169 ms: Couldn't connect to server

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]

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

[metascroy]

@claude review this code

[claude]

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

---

Code Review — DFlash (Qwen3-4B), branch dflash-qwen3-4b @ fb44b86

Big step forward since the 2026-08-03 pass. I re-read the whole branch. The five structural landing criteria are now largely met; the remaining issues cluster around the draft KV cache path (criterion 3), which is implemented but off by default, unvalidated by any runnable check, and shipped alongside two stale scratch files whose claimed "verification" no longer applies to the code.

Landing criteria — status

#	Item	Status
1	Dynamic block_len w/ generation + context bounds	✅ compute_block_len + Dim("block_len", min=2, max=block_size); bs==1 → target-only step
2	Reuse HF reference module	✅ DFlashQwen3Attention(Qwen3Attention) reuses Qwen3RMSNorm/MLP/RotaryEmbedding + ALL_ATTENTION_FUNCTIONS dispatch; embed_scale now loaded; eager-parity test added
3	Persistent draft KV cache	🟡 implemented (DFlashDraftKVCache, --cached) but opt-in, unvalidated, docs contradict code — see below
4	One generic Qwen-first path	✅ Gemma impl removed (7261f77), single path
5	Explicit hidden tapping	✅ renamed --tap-layers; sliding-window combo rejected explicitly
6	Validation scripts	✅ lossless fails on nonzero exit + mismatch; target check exercises per-token positions across seq lens; synthetic acceptance-loop tests added

🔴 Blocking

1. check_dflash_draft.py passes 3 inputs to a 2-input exported program. The uncached export (export_dflash_draft.py:169) traces model with inputs (tokens, target_hidden) — positions are computed inside DFlashDraftModel.forward (dflash_draft_model.py:305). But check_dflash_draft.py:29 calls method.execute([tokens, target_hidden, position_ids]) with a third position_ids tensor the forward never accepted. run_dflash.py:248 correctly calls draft.execute([draft_input, hidden]) (2 args), so this check is inconsistent with both the export and the real runner and will fail against a genuinely-exported draft. Drop the position_ids arg. Fix →

2. Committed scratch files are stale and reference a cache API that no longer exists. Both are tracked (git ls-files confirms) and both are self-described as "Not part of the PR — a throwaway diagnostic":

• scratch_draft_cache_probe.py calls cache.write_context(), cache.advance_context(), cache.ctx_len, cache.write_block(), and cache.valid_mask(block_len=8) — none of these exist on the shipped DFlashDraftKVCache (which exposes write(), valid_mask(valid_len, device), valid_len_after(), reset()).
• scratch_draft_cache_equiv.py:70-73 calls model(tokens, ctx, cache=..., cache_position=..., new_ctx_len=new_ctx_len) — but DFlashDraftModel.forward takes no new_ctx_len kwarg (it derives it from target_hidden.shape[1]). This raises TypeError on run.

So the dflash_draft_model.py:210 claim that the .item()+narrow() cache path was "verified directly via a standalone probe before wiring this in" is not backed by any runnable code on this branch — the probe it refers to tests a different, older cache design. Please git rm both scratch files (the reviewer's earlier "remove experiment files not part of the supported workflow" applies here). Fix →

🟡 Should fix

3. The draft KV cache (criterion 3) has no runnable validation and is off by default. --cached is opt-in in both export_dflash_draft.py and run_dflash.py, and check_dflash_lossless.py runs run_dflash.py without --cached — so the only lossless guarantee covers the uncached (full-reprojection) path. The eager-equivalence check that was meant to validate the cached path (scratch_draft_cache_equiv.py) is broken (#2). Since the persistent cache is the actual point-3 deliverable and the perf argument for it, it needs a real, committed test: either a check_dflash_cached_lossless.py that runs run_dflash.py --cached and compares to baseline, or a fixed eager cached-vs-uncached parity test under tests/. Without it, the cache path can silently regress.

4. Contradictory / duplicated comments in dflash_draft_model.py cached path. The block comment at :149-155 says the design attends "over the FULL fixed-size cache buffer with a validity mask -- never narrow() it down," but the code at :217-221 does exactly the opposite — valid_len = ...item() then full_k.narrow(2, 0, valid_len). The comment at :199-201 ("Write new context K/V into the persistent cache…") is also a stale duplicate of :189-191, describing a write that already happened at :191. And DFlashDraftKVCache.valid_mask() / its docstring (dflash_draft_cache.py:62-74) describe the mask-based approach that the model no longer uses — so valid_mask() is now dead code. Reconcile the comments with the narrow-based implementation actually shipped, and either wire up or remove valid_mask(). This directly obscures a correctness-sensitive path.

5. compute_block_len() is tested but not used by the real loop. run_dflash.py:204 reimplements the bound inline (bs = min(block_size, ...)) and :208-209 reimplements the draft_room bound, rather than calling compute_block_len(). So the synthetic tests in check_dflash_acceptance_loop.py exercise a helper the production loop doesn't call — the two can drift. Call compute_block_len(...) from the loop so the tests cover the real path. Fix →

🟢 Minor / nits

• dflash_draft_cache.py:59 write() docstring says "Caller masks the unwritten tail via valid_mask()" — stale; the caller narrows, doesn't mask.
• check_dflash_draft.py:22 hardcodes hidden_size=12800 (5 tapped layers × 2560). Fine for a smoke check, but a comment tying it to --tap-layers/target_layer_ids would prevent confusion if the tap set changes.
• export_dflash_draft.py:118-119 docstring references scratch_draft_cache_export.py, which isn't in the tree — dangling reference; drop it once Re-sync with internal repository #2 is done.
• run_dflash.py:305 / :311 the .float() on new_hidden and the two context-accumulation branches are correct but subtle; the existing comments cover it well — no change needed.

Summary

Criteria 1, 2, 4, 5, 6 are in good shape. The blockers are small and mechanical: fix the 3-arg draft check (#1) and remove the two stale scratch files (#2). The most substantive remaining gap is that the persistent draft cache (criterion 3) ships without a runnable lossless/parity check and with comments that contradict the shipped narrow() implementation (#3, #4) — worth closing before landing, since that path is the whole performance rationale for the cache.
· branch dflash-qwen3-4b

[cthotti]

Hi, Thanks for the thorough review @metascroy, and sorry for the slow write-up. All six points are addressed below:

1. Dynamic block_len

The proposal dim now exports as a bounded torch.export.Dim (1 <= block_len <= checkpoint_block_size), and the runner picks the length from min(checkpoint_block_size, max_new_tokens - num_generated, max_seq_len - position). Positions are derived inside the draft wrapper from the input shapes, so block_len and ctx_len can stay together. This needed an operator.neg handler in backends/mlx/ops.py, and the bounded dims introduce negations in the position math that didn't exist with a constant. block_len == 1 lowers cleanly, so no target-only path was necessary. I confirmed with non-multiple max-new-tokens that generation stops exactly on budget and the final 1-token round skips the draft.

2. HF reference module

I swapped the hand-rolled stack for adapted Qwen3Attention/Qwen3DecoderLayer through ALL_ATTENTION_FUNCTIONS["mlx"]. embed_scale is loaded and applied, and the sliding-window path no longer forces a causal mask (Qwen3 is noncausal). A thin wrapper copies the target embed_tokens/lm_head in so the draft stays self-contained. The eager parity test against real Qwen3Attention is bit-exact and I also checked full DFlashDraftModel.forward(), MLX lowering, and weight loading against the z-lab checkpoint (names line up, no remapping).

3. Draft KV cache

DFlashDraftKVCache follows the TorchExportableModuleWithStaticCache pattern (mutable-buffer tensors, cache_position threaded through), with append/crop made explicit, valid_len_after() for the crop and valid_mask() gating populated slots. But two bugs turned up in the process.

• First, the draft cache had no bound independent of the target's max-seq-len and could silently overrun, so now bounded by --draft-max-ctx-len with a hard assert (a too-small bound shows rounds: 45 vs. the correct rounds: 26, tau: 5.73).
• Second, my initial masked-attention version was slower than uncached at every length because cost tracked buffer size, not real context. But switching to narrow()-based slicing (export-safe via .item() and torch._check(), same as custom_sdpa) came out 1.68x faster at max_seq_len=4096.

Eager parity is 7e-07 max diff, and draft cost now stays flat as context grows.

4. Consolidation, validation, cleanup

Gemma4-31B (exporters, Python/C++ drivers, CMakeLists) is now on its own dflash-gemma4-31b branch with nothing lost, leaving a single generic path: export_dflash_draft.py and run_dflash.py under backends/mlx/examples/llm/, Qwen3 specifics in dflash_qwen3_adapter.py. --dflash-layers to --tap-layers, with tap-layers + sliding-window now rejected on export. On the scripts:

• The smoke check passes one position per token across multiple lengths
• The losslessness check fails on subprocess errors and mismatches instead of passing empty
• And check_dflash_acceptance_loop.py covers first-token rejection, full acceptance with bonus, EOS, and a shrunk final budget.

I also removed scratch files, a stray installer, and trailing whitespace via lintrunner.

One aside, out of scope:

While checking batched verification I root-caused a pre-existing backend issue.get_qmv_batch_limit() flips the quantized matmul kernel from qmv to qmm at batch size 6, and the accumulation-order difference (~0.18–0.31 logits) is enough to flip argmax on close calls. It reproduces with no DFlash code in the path, so it's a backend property, not ours. --block-size 5 keeps verification under the threshold and holds exact losslessness over 200 tokens. Flagging this since it came straight out of this work.

Let me know if there's anything else you'd like changed. Happy to keep iterating.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 7m 55s —— View job

---

Code Review — DFlash (Qwen3-4B), branch dflash-qwen3-4b

I re-read the branch against origin/main. Big structural change since the 2026-08-10 pass: the DFlash code is now a single generic package under backends/mlx/examples/llm/dflash/ (arch.py, adapters.py, model.py, cache.py, export.py, run.py), the hand-rolled draft stack is gone in favor of subclassing the checkpoint's own HF attention/decoder-layer via resolve_arch, and operator.neg lowering was added. This is a genuinely cleaner design — the arch/adapters split isolates the one truly per-checkpoint piece (weight loading) well, and model.py's DFlashAttention.adapt() reuse is elegant.

Against your six landing criteria the structural items (1, 2, 4, 5) are in good shape. But re-reviewing surfaced one blocking correctness bug and a significant test-coverage regression introduced by the refactor.

🔴 Blocking

1. The runner's default (uncached) path can't run any .pte this exporter produces. dflash/export.py only ever exports a cached draft — to_edge_transform_and_lower({"target": ..., "draft": exported}) at export.py:234-235, where exported is the 3-input DFlashCachedDraftModel(tokens, new_ctx, cache_position) (export.py:169-195). But --cached in run.py:105-110 is action="store_true" with default False, and the uncached branch calls draft.execute([draft_input, hidden]) — 2 args (run.py:271) — against a method that always expects 3. So the default invocation (python -m …dflash.run --pte x.pte …, no --cached) fails at runtime; only --cached works. The CI job passes --cached explicitly (.github/workflows/mlx.yml), so this is entirely hidden. Either make --cached the default (and drop/guard the uncached branch), or also export an uncached draft method. As written the out-of-the-box path is broken. Fix →

🟡 Should fix

2. The refactor deleted all the in-repo validation added for criterion 6. Commit 795c37dc removed every examples/models/qwen3/tests/check_dflash_*.py: check_dflash_lossless.py, check_dflash_acceptance_loop.py (the synthetic first-token-rejection / full-accept-with-bonus / EOS / budget-<-block tests), check_dflash_draft_parity.py, check_dflash_cached_parity.py, check_dflash_target.py. Nothing replaced them in dflash/. The new test-mlx-dflash CI job is a nice end-to-end smoke test, but it does not check losslessness (no token-for-token comparison against baseline greedy decoding — it only asserts "Paris" appears and 1.0 < tau <= n_draft+1), and the pure-Python loop helpers (first_mismatch, compute_block_len, truncate_at_eos) and the eager-vs-HF parity check now have no test at all. Losslessness — the headline correctness claim of this PR — regressed from "manually verified + scripted" to "not asserted anywhere." Please re-add at least (a) a lossless baseline-vs-DFlash comparison to the CI job, and (b) a small synthetic unit test for the acceptance loop that runs without hardware/HF downloads.

3. --cached vs the exporter make the whole cache path the only real path, yet the model still carries a full uncached branch. DFlashAttention.forward (model.py:143-159) and DFlashDraftModel.forward (model.py:308-312) keep an uncached branch described as "reference behavior," but per #1 it's never exported and (after this refactor) never tested. Either wire it into an exported method + parity test, or drop it so model.py doesn't carry a second attention/position path that can silently diverge from the shipped one.

4. Claimed "tap-layers + sliding-window now rejected" doesn't match the code. Your write-up says the combination is rejected on export, but create_hf_exportable (exportable.py:143-153) actually supports it — when sliding_window is not None and tap_layers is not None it builds TorchExportableModuleWithHybridCacheAndHidden. Supporting it is arguably better than rejecting, but the hybrid+tap wrapper is untested here and contradicts the stated behavior. Please reconcile: either test that path or explicitly reject as claimed.

🟢 Minor / nits

• dflash/__init__.py:1-2 has only the Copyright line — missing the two This source code is licensed under the BSD-style license … lines every other file in the package carries.
• model.py:163 the comment on the cached K/V narrow is a single long unwrapped line; wrap to match the file's style.
• run.py:65 truncate_at_eos: if eos_id in new_tokens is safe when eos_id is None (None in list → False), but a one-line note that get_eos_token_id may return None would save a future reader the double-take.
• .github/workflows/mlx.yml tau parse: grep -oE "[0-9]+\.[0-9]+" on the "avg emitted/round" line is fine today, but it silently takes the first float on the line — worth anchoring to the value to avoid a future format change matching the wrong number.
• arch.py / model.py lru_cache on class factories: correct and keeps state-dict layout stable, but note these caches are process-global keyed by model_type — fine for the export tool, just be aware if this is ever imported into a long-lived multi-model process.

Landing-criteria status

#	Item	Status
1	Dynamic block_len	✅ Dim("block_len", min=2, max=block_size); compute_block_len now called by the real loop (run.py:231)
2	Reuse HF reference module	✅ DFlashAttention(arch.attention_cls) + ALL_ATTENTION_FUNCTIONS["mlx"], embed_scale loaded, unsupported-layer rejection
3	Persistent draft KV cache	🟡 implemented and now the only exported path — but uncached default runner broken (#1), no lossless/parity test (#2)
4	One generic Qwen-first path	✅ Gemma removed; single dflash/ package with thin adapter
5	Explicit hidden tapping	✅ --tap-layers; but tap+sliding is supported, not rejected as claimed (#4)
6	Reliable validation	❌ regressed — in-repo checks deleted, CI doesn't assert losslessness (#2)

Summary

The generic redesign is a clear improvement and criteria 1, 2, 4 are solid. Two things to resolve before landing: the default (uncached) runner path is broken against the only .pte the exporter produces (#1), and the refactor dropped the losslessness/acceptance-loop/parity tests so the PR's core correctness claim is no longer verified anywhere in-repo (#2). Fixing those two makes this landable.
· branch dflash-qwen3-4b

[metascroy]

Thanks for the PR @cthotti! I just cleaned up the logic a bit to make it more generic, rather than Qwen3 specific. Once CI passes, we can merge.

Note I added an n_draft parameter to the dflash run.py command. Setting it to something small like 3 or 4 significantly speeds things up for me compared to using 16.