From ed80f41c5023f84f2b3079b0c1459bf5cdf91f05 Mon Sep 17 00:00:00 2001
From: cthotti <cthotti@gmail.com>
Date: Sat, 4 Jul 2026 14:15:50 +0530
Subject: [PATCH 1/4] DFlash Phase 2: draft model definition, HF weight
 loading, and .pte export for Qwen3-4B

---
 .../mlx/examples/llm/dflash_draft_model.py    | 227 ++++++++++++++++++
 examples/models/qwen3/export_dflash_draft.py  |  67 ++++++
 .../qwen3/tests/test_dflash_draft_pte.py      |  19 ++
 3 files changed, 313 insertions(+)
 create mode 100644 backends/mlx/examples/llm/dflash_draft_model.py
 create mode 100644 examples/models/qwen3/export_dflash_draft.py
 create mode 100644 examples/models/qwen3/tests/test_dflash_draft_pte.py

diff --git a/backends/mlx/examples/llm/dflash_draft_model.py b/backends/mlx/examples/llm/dflash_draft_model.py
new file mode 100644
index 00000000000..a573e64142b
--- /dev/null
+++ b/backends/mlx/examples/llm/dflash_draft_model.py
@@ -0,0 +1,227 @@
+"""PyTorch DFlash draft model, structured for ExecuTorch export.
+
+Model-agnostic: the same code exports a valid draft .pte for any standard-
+attention target (Qwen3, Gemma-4, Llama-3.1) by reading a DFlashConfig loaded
+from the z-lab draft checkpoint. Per-model differences — RoPE base/scaling,
+sliding-window layers, final-logit softcap, embedding scale — are config values,
+so they resolve at trace time to one model-specific graph. The universal
+branches add no ops to the exported program.
+
+Two deliberate deviations from the reference forward, both for ET (design doc
+"Option A", self-contained draft):
+  - embed_tokens / lm_head are owned here (filled from the target at export)
+    instead of referenced live off the target module.
+  - forward returns draft logits (norm -> lm_head -> [:, 1:]) rather than the
+    bare normed hidden state; the reference does lm_head + logits_start=1 in its
+    generate loop.
+
+References:
+  z-lab/dflash  dflash/model_mlx.py  — universal MLX reference (Qwen3 / Qwen3.5 /
+      Gemma-4); source of the sliding-window, softcap, single-rope, QK-norm design.
+  z-lab/dflash  dflash/model.py       — PyTorch reference; exact weight layout for
+      load_state_dict.
+  transformers RoPE init               — config-driven inv_freq below.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+from torch import nn
+
+
+@dataclass
+class DFlashConfig:
+    hidden_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    intermediate_size: int
+    vocab_size: int
+    rms_norm_eps: float
+    rope_theta: float
+    max_position_embeddings: int
+    target_layer_ids: Tuple[int, ...]
+    block_size: int = 16
+    mask_token_id: int = 0
+    rope_scaling: Optional[Dict[str, Any]] = None
+    layer_types: Tuple[str, ...] = field(default_factory=tuple)
+    sliding_window: Optional[int] = None
+    final_logit_softcapping: Optional[float] = None
+    embed_scale: float = 1.0  # 1.0 for Qwen3/Llama; sqrt(hidden_size) for Gemma
+
+
+def _rope_inv_freq(config: DFlashConfig) -> torch.Tensor:
+    # Covers the two rope types these drafts use: 'default' (Qwen3, Gemma-4) and
+    # 'linear' position-scaling. YaRN/longrope aren't handled — no current z-lab
+    # draft uses them.
+    dim = config.head_dim
+    inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    scaling = config.rope_scaling or {}
+    if scaling.get("rope_type", scaling.get("type")) == "linear":
+        inv_freq = inv_freq / float(scaling["factor"])
+    return inv_freq
+
+
+class DFlashRotaryEmbedding(nn.Module):
+    def __init__(self, config: DFlashConfig):
+        super().__init__()
+        self.register_buffer("inv_freq", _rope_inv_freq(config), persistent=False)
+
+    def forward(self, position_ids: torch.Tensor):
+        inv = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        pos = position_ids[:, None, :].float()
+        freqs = (inv @ pos).transpose(1, 2)      # [B, S, head_dim/2]
+        emb = torch.cat((freqs, freqs), dim=-1)  # [B, S, head_dim]
+        return emb.cos(), emb.sin()
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    # q rotates over its own (last q_len) positions; k rotates over its full length.
+    q_len = q.shape[-2]
+    cq, sq = cos[:, None, -q_len:, :], sin[:, None, -q_len:, :]
+    ck, sk = cos[:, None, :, :], sin[:, None, :, :]
+    return (q * cq) + (rotate_half(q) * sq), (k * ck) + (rotate_half(k) * sk)
+
+
+class DFlashRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return self.weight * x.to(dtype)
+
+
+class DFlashMLP(nn.Module):
+    def __init__(self, hidden_size: int, intermediate_size: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DFlashAttention(nn.Module):
+    def __init__(self, config: DFlashConfig, layer_idx: int):
+        super().__init__()
+        h, hd = config.hidden_size, config.head_dim
+        self.n_heads = config.num_attention_heads
+        self.n_kv = config.num_key_value_heads
+        self.head_dim = hd
+        self.scaling = hd ** -0.5
+        self.n_rep = self.n_heads // self.n_kv
+        lt = config.layer_types
+        self.is_sliding = bool(lt) and lt[layer_idx] == "sliding_attention"
+        self.sliding_window = config.sliding_window if self.is_sliding else None
+        self.q_proj = nn.Linear(h, self.n_heads * hd, bias=False)
+        self.k_proj = nn.Linear(h, self.n_kv * hd, bias=False)
+        self.v_proj = nn.Linear(h, self.n_kv * hd, bias=False)
+        self.o_proj = nn.Linear(self.n_heads * hd, h, bias=False)
+        self.q_norm = DFlashRMSNorm(hd, config.rms_norm_eps)
+        self.k_norm = DFlashRMSNorm(hd, config.rms_norm_eps)
+
+    def forward(self, x, x_ctx, cos, sin):
+        B, L, _ = x.shape
+        S = x_ctx.shape[1]
+        q = self.q_norm(self.q_proj(x).view(B, L, self.n_heads, self.head_dim)).transpose(1, 2)
+        k = torch.cat([self.k_proj(x_ctx), self.k_proj(x)], dim=1).view(B, S + L, self.n_kv, self.head_dim)
+        v = torch.cat([self.v_proj(x_ctx), self.v_proj(x)], dim=1).view(B, S + L, self.n_kv, self.head_dim)
+        k = self.k_norm(k).transpose(1, 2)
+        v = v.transpose(1, 2)
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+        if self.n_rep > 1:
+            k = k.repeat_interleave(self.n_rep, dim=1)
+            v = v.repeat_interleave(self.n_rep, dim=1)
+        mask = self._sliding_mask(L, S, q.device, q.dtype) if self.is_sliding else None
+        out = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=mask, is_causal=False, scale=self.scaling)
+        return self.o_proj(out.transpose(1, 2).reshape(B, L, -1))
+
+    def _sliding_mask(self, L, S, device, dtype):
+        total = S + L
+        q_pos = torch.arange(S, total, device=device)[:, None]
+        k_pos = torch.arange(total, device=device)[None, :]
+        allowed = (k_pos <= q_pos) & (k_pos > q_pos - self.sliding_window)
+        return torch.where(allowed, 0.0, float("-inf")).to(dtype)[None, None]
+
+
+class DFlashDecoderLayer(nn.Module):
+    def __init__(self, config: DFlashConfig, layer_idx: int):
+        super().__init__()
+        self.self_attn = DFlashAttention(config, layer_idx)
+        self.mlp = DFlashMLP(config.hidden_size, config.intermediate_size)
+        self.input_layernorm = DFlashRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = DFlashRMSNorm(config.hidden_size, config.rms_norm_eps)
+
+    def forward(self, x, x_ctx, cos, sin):
+        x = x + self.self_attn(self.input_layernorm(x), x_ctx, cos, sin)
+        return x + self.mlp(self.post_attention_layernorm(x))
+
+
+class DFlashDraftModel(nn.Module):
+    def __init__(self, config: DFlashConfig):
+        super().__init__()
+        self.config = config
+        concat_dim = len(config.target_layer_ids) * config.hidden_size
+        self.fc = nn.Linear(concat_dim, config.hidden_size, bias=False)
+        self.hidden_norm = DFlashRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.layers = nn.ModuleList(
+            [DFlashDecoderLayer(config, i) for i in range(config.num_hidden_layers)])
+        self.norm = DFlashRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.rotary_emb = DFlashRotaryEmbedding(config)
+        # Option A: owned copies, filled from the target at export time.
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def forward(self, tokens, target_hidden, position_ids):
+        h = self.embed_tokens(tokens) * self.config.embed_scale
+        h_ctx = self.hidden_norm(self.fc(target_hidden))
+        cos, sin = self.rotary_emb(position_ids)
+        for layer in self.layers:
+            h = layer(h, h_ctx, cos, sin)
+        h = self.norm(h)
+        logits = self.lm_head(h[:, 1:, :])   # logits_start=1: drop the known first token
+        cap = self.config.final_logit_softcapping
+        if cap is not None:
+            logits = torch.tanh(logits / cap) * cap
+        return logits
+
+def load_dflash_config(checkpoint_dir) -> "DFlashConfig":
+    """Build a DFlashConfig from a z-lab DFlash checkpoint's config.json."""
+    import json
+    from pathlib import Path
+
+    cfg = json.loads((Path(checkpoint_dir) / "config.json").read_text())
+    dcfg = cfg["dflash_config"]
+    return DFlashConfig(
+        hidden_size=cfg["hidden_size"],
+        num_hidden_layers=cfg["num_hidden_layers"],
+        num_attention_heads=cfg["num_attention_heads"],
+        num_key_value_heads=cfg["num_key_value_heads"],
+        head_dim=cfg["head_dim"],
+        intermediate_size=cfg["intermediate_size"],
+        vocab_size=cfg["vocab_size"],
+        rms_norm_eps=cfg["rms_norm_eps"],
+        rope_theta=cfg["rope_theta"],
+        max_position_embeddings=cfg["max_position_embeddings"],
+        target_layer_ids=tuple(dcfg["target_layer_ids"]),
+        block_size=cfg["block_size"],
+        mask_token_id=dcfg["mask_token_id"],
+        rope_scaling=cfg.get("rope_scaling"),
+        layer_types=tuple(cfg.get("layer_types") or ["full_attention"] * cfg["num_hidden_layers"]),
+        sliding_window=cfg.get("sliding_window"),
+        final_logit_softcapping=cfg.get("final_logit_softcapping"),
+    )
diff --git a/examples/models/qwen3/export_dflash_draft.py b/examples/models/qwen3/export_dflash_draft.py
new file mode 100644
index 00000000000..838ab14b6a5
--- /dev/null
+++ b/examples/models/qwen3/export_dflash_draft.py
@@ -0,0 +1,67 @@
+import argparse
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+from transformers import AutoModelForCausalLM
+
+from executorch.backends.mlx.examples.llm.dflash_draft_model import DFlashDraftModel, load_dflash_config
+
+
+def load_draft_model(draft_id: str, target_state_dict: dict) -> DFlashDraftModel:
+    path = Path(snapshot_download(draft_id, allow_patterns=["*.safetensors", "*.json"]))
+    config = load_dflash_config(path)
+    model = DFlashDraftModel(config)
+
+    draft_weights = {}
+    for f in path.glob("*.safetensors"):
+        draft_weights.update(load_file(str(f)))
+
+    missing, unexpected = model.load_state_dict(draft_weights, strict=False)
+    assert not unexpected, f"Unexpected draft checkpoint keys: {unexpected}"
+    still_missing = [k for k in missing if not k.startswith(("embed_tokens.", "lm_head."))]
+    assert not still_missing, f"Missing draft checkpoint keys: {still_missing}"
+
+    model.embed_tokens.weight.data.copy_(target_state_dict["model.embed_tokens.weight"])
+    lm_head_key = "lm_head.weight" if "lm_head.weight" in target_state_dict else "model.embed_tokens.weight"
+    model.lm_head.weight.data.copy_(target_state_dict[lm_head_key])
+    return model
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--target-model", default="Qwen/Qwen3-4B")
+    parser.add_argument("--draft-model", default="z-lab/Qwen3-4B-DFlash-b16")
+    parser.add_argument("--output", default="qwen3_4b_dflash_draft.pte")
+    parser.add_argument("--block-size", type=int, default=16)
+    parser.add_argument("--ctx-len", type=int, default=8)
+    args = parser.parse_args()
+
+    target = AutoModelForCausalLM.from_pretrained(args.target_model, dtype="auto")
+    model = load_draft_model(args.draft_model, target.state_dict())
+    model.eval()
+    model = model.float()
+    del target
+
+    block_size, ctx_len = args.block_size, args.ctx_len
+    hidden_size = model.fc.in_features
+    tokens = torch.randint(0, 1000, (1, block_size), dtype=torch.long)
+    target_hidden = torch.randn(1, ctx_len, hidden_size)
+    position_ids = torch.arange(ctx_len + block_size).unsqueeze(0)
+
+    exported = torch.export.export(model, (tokens, target_hidden, position_ids))
+
+    from executorch.exir import to_edge_transform_and_lower
+    from executorch.backends.mlx.partitioner import MLXPartitioner
+
+    edge = to_edge_transform_and_lower(exported, partitioner=[MLXPartitioner()])
+    et_program = edge.to_executorch()
+
+    with open(args.output, "wb") as f:
+        f.write(et_program.buffer)
+    print(f"Saved draft model to: {args.output}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/models/qwen3/tests/test_dflash_draft_pte.py b/examples/models/qwen3/tests/test_dflash_draft_pte.py
new file mode 100644
index 00000000000..374bb0e2aa0
--- /dev/null
+++ b/examples/models/qwen3/tests/test_dflash_draft_pte.py
@@ -0,0 +1,19 @@
+import sys
+import torch
+from executorch.runtime import Runtime, Verification
+
+pte_path = sys.argv[1] if len(sys.argv) > 1 else "qwen3_4b_dflash_draft.pte"
+et_runtime = Runtime.get()
+program = et_runtime.load_program(pte_path, verification=Verification.Minimal)
+method = program.load_method("forward")
+
+# Must match the exact static shapes used at export time: ctx_len=8, block_size=16
+block_size, ctx_len, hidden_size, vocab_size = 16, 8, 12800, 151936
+tokens = torch.randint(0, 1000, (1, block_size), dtype=torch.long)
+target_hidden = torch.randn(1, ctx_len, hidden_size)
+position_ids = torch.arange(ctx_len + block_size).unsqueeze(0).long()
+
+(draft_logits,) = method.execute([tokens, target_hidden, position_ids])
+assert draft_logits.shape == (1, block_size - 1, vocab_size), draft_logits.shape
+assert not torch.isnan(draft_logits).any() and not torch.isinf(draft_logits).any()
+print(f"OK: draft_logits {tuple(draft_logits.shape)}")

From e0420a374629ae644830e0a888ffd7d72232acf9 Mon Sep 17 00:00:00 2001
From: cthotti <cthotti@gmail.com>
Date: Sun, 5 Jul 2026 21:15:14 +0530
Subject: [PATCH 2/4] DFlash: consolidate tests

---
 .gitignore                                    |   1 +
 Makefile                                      |   9 +
 .../mlx/examples/llm/dflash_draft_model.py    |  10 +-
 backends/mlx/examples/llm/export_llm_hf.py    |  37 +++-
 examples/models/qwen3/export_dflash_draft.py  |  32 +++-
 .../qwen3/mlx_source_transformations.py       | 110 ++++++++++++
 examples/models/qwen3/run_baseline.py         |  69 ++++++++
 examples/models/qwen3/run_dflash.py           | 162 ++++++++++++++++++
 .../models/qwen3/tests/test_dflash_draft.py   |  28 +++
 .../qwen3/tests/test_dflash_draft_dynamic.py  |  18 ++
 .../qwen3/tests/test_dflash_draft_eager.py    |  34 ++++
 .../qwen3/tests/test_dflash_draft_forward.py  |  46 +++++
 .../tests/test_dflash_draft_load_weights.py   |  50 ++++++
 .../models/qwen3/tests/test_dflash_export.py  |  32 ++++
 .../qwen3/tests/test_dflash_lossless.py       |  32 ++++
 .../models/qwen3/tests/test_dflash_target.py  |  32 ++++
 16 files changed, 695 insertions(+), 7 deletions(-)
 create mode 100644 examples/models/qwen3/mlx_source_transformations.py
 create mode 100644 examples/models/qwen3/run_baseline.py
 create mode 100644 examples/models/qwen3/run_dflash.py
 create mode 100644 examples/models/qwen3/tests/test_dflash_draft.py
 create mode 100644 examples/models/qwen3/tests/test_dflash_draft_dynamic.py
 create mode 100644 examples/models/qwen3/tests/test_dflash_draft_eager.py
 create mode 100644 examples/models/qwen3/tests/test_dflash_draft_forward.py
 create mode 100644 examples/models/qwen3/tests/test_dflash_draft_load_weights.py
 create mode 100644 examples/models/qwen3/tests/test_dflash_export.py
 create mode 100644 examples/models/qwen3/tests/test_dflash_lossless.py
 create mode 100644 examples/models/qwen3/tests/test_dflash_target.py

diff --git a/.gitignore b/.gitignore
index ee206e23d94..10b048c17b5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -69,6 +69,7 @@ xcuserdata/
 /src/executorch/include/
 /src/executorch/share/
 /src/executorch/version.py
+/dflash_benchmarks.md
 *_etdump
 
 # Android
diff --git a/Makefile b/Makefile
index 969b53644cd..26167aeb2b2 100644
--- a/Makefile
+++ b/Makefile
@@ -484,3 +484,12 @@ clean:
 	rm -rf cmake-out \
 	       extension/llm/tokenizers/build \
 	       extension/llm/tokenizers/pytorch_tokenizers.egg-info
+
+qwen3_dflash-mlx:
+	@echo "==> Building and installing ExecuTorch with MLX..."
+	cmake --workflow --preset mlx-release
+	@echo "==> Building Qwen3 DFlash speculative decoding runner with MLX..."
+	cd examples/models/qwen3 && cmake --workflow --preset qwen3-dflash-mlx
+	@echo ""
+	@echo "✓ Build complete!"
+	@echo "  Runner: cmake-out/examples/models/qwen3/qwen3_dflash_runner"
diff --git a/backends/mlx/examples/llm/dflash_draft_model.py b/backends/mlx/examples/llm/dflash_draft_model.py
index a573e64142b..913f89e2d57 100644
--- a/backends/mlx/examples/llm/dflash_draft_model.py
+++ b/backends/mlx/examples/llm/dflash_draft_model.py
@@ -81,6 +81,12 @@ def rotate_half(x: torch.Tensor) -> torch.Tensor:
     x1, x2 = x.chunk(2, dim=-1)
     return torch.cat((-x2, x1), dim=-1)
 
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    if n_rep == 1:
+        return x
+    b, h, s, d = x.shape
+    x = x[:, :, None, :, :].expand(b, h, n_rep, s, d)
+    return x.reshape(b, h * n_rep, s, d)
 
 def apply_rotary_pos_emb(q, k, cos, sin):
     # q rotates over its own (last q_len) positions; k rotates over its full length.
@@ -143,8 +149,8 @@ def forward(self, x, x_ctx, cos, sin):
         v = v.transpose(1, 2)
         q, k = apply_rotary_pos_emb(q, k, cos, sin)
         if self.n_rep > 1:
-            k = k.repeat_interleave(self.n_rep, dim=1)
-            v = v.repeat_interleave(self.n_rep, dim=1)
+            k = repeat_kv(k, self.n_rep)
+            v = repeat_kv(v, self.n_rep)
         mask = self._sliding_mask(L, S, q.device, q.dtype) if self.is_sliding else None
         out = torch.nn.functional.scaled_dot_product_attention(
             q, k, v, attn_mask=mask, is_causal=False, scale=self.scaling)
diff --git a/backends/mlx/examples/llm/export_llm_hf.py b/backends/mlx/examples/llm/export_llm_hf.py
index fe6b8094f6b..45fd8bfad28 100644
--- a/backends/mlx/examples/llm/export_llm_hf.py
+++ b/backends/mlx/examples/llm/export_llm_hf.py
@@ -137,6 +137,7 @@ def _export_with_custom_components(
     no_tie_word_embeddings: bool = False,
     qlinear_group_size: Optional[int] = None,
     qembedding_group_size: Optional[int] = None,
+    dflash_layers: Optional[list[int]] = None,
 ) -> None:
     """
     Export using direct HF model with custom MLX components.
@@ -219,6 +220,24 @@ def _export_with_custom_components(
             batch_size=1,
             max_cache_len=effective_cache_len,
         )
+    elif dflash_layers is not None:
+        # Qwen3-specific for now
+        # Generalize the import if/when another model needs DFlash tapping.
+        #
+        # Stateless (no persistent KV cache), NOT TorchExportableModuleWithStaticCacheAndHidden:
+        # DFlash's speculative-decode loop re-verifies overlapping/non-contiguous token
+        # ranges every round, which corrupts a persistent StaticCache (confirmed at the
+        # eager PyTorch level -- see StatelessQwen3WithHidden's docstring). The driver
+        # always passes the full accumulated sequence instead of relying on a cache.
+        from executorch.examples.models.qwen3.mlx_source_transformations import (
+            StatelessQwen3WithHidden,
+        )
+
+        logger.info(f"Creating stateless DFlash hidden-state-tapping wrapper, layers={dflash_layers}")
+        exportable = StatelessQwen3WithHidden(
+            model=model,
+            layer_ids=dflash_layers,
+        )
     else:
         logger.info("Creating TorchExportableModuleWithStaticCache wrapper...")
         exportable = TorchExportableModuleWithStaticCache(
@@ -227,7 +246,7 @@ def _export_with_custom_components(
             max_cache_len=effective_cache_len,
         )
 
-    if use_custom_kv_cache:
+    if use_custom_kv_cache and dflash_layers is None:
         from executorch.backends.mlx.llm.source_transformation import (
             replace_hf_cache_with_mlx,
         )
@@ -335,6 +354,7 @@ def export_llama_hf(
     no_tie_word_embeddings: bool = False,
     qlinear_group_size: Optional[int] = None,
     qembedding_group_size: Optional[int] = None,
+    dflash_layers: Optional[list[int]] = None,
 ) -> None:
     """
     Export a HuggingFace Llama model to ExecuTorch with MLX backend.
@@ -349,10 +369,10 @@ def export_llama_hf(
         use_custom_sdpa: Use MLX custom SDPA (mlx::custom_sdpa)
         use_custom_kv_cache: Use MLX custom KV cache (mlx::kv_cache_update)
     """
-    if use_custom_sdpa or use_custom_kv_cache:
+    if use_custom_sdpa or use_custom_kv_cache or dflash_layers is not None:
         logger.info(
             f"Using custom components: sdpa={use_custom_sdpa}, "
-            f"kv_cache={use_custom_kv_cache}"
+            f"kv_cache={use_custom_kv_cache}, dflash_layers={dflash_layers}"
         )
         _export_with_custom_components(
             model_id=model_id,
@@ -367,6 +387,7 @@ def export_llama_hf(
             no_tie_word_embeddings=no_tie_word_embeddings,
             qlinear_group_size=qlinear_group_size,
             qembedding_group_size=qembedding_group_size,
+            dflash_layers=dflash_layers,
         )
     else:
         logger.info("Using optimum-executorch pipeline (no custom components)")
@@ -434,8 +455,17 @@ def main():
         default=False,
         help="Use MLX custom KV cache (mlx::kv_cache_update)",
     )
+    parser.add_argument(
+        "--dflash-layers",
+        type=str,
+        default=None,
+        help="Comma-separated layer indices to tap for DFlash hidden-state output, e.g. '2,18,33'",
+    )
 
     args = parser.parse_args()
+    dflash_layers = (
+        [int(x) for x in args.dflash_layers.split(",")] if args.dflash_layers else None
+    )
 
     export_llama_hf(
         model_id=args.model_id,
@@ -450,6 +480,7 @@ def main():
         no_tie_word_embeddings=args.no_tie_word_embeddings,
         qlinear_group_size=args.qlinear_group_size,
         qembedding_group_size=args.qembedding_group_size,
+        dflash_layers=dflash_layers,
     )
 
 
diff --git a/examples/models/qwen3/export_dflash_draft.py b/examples/models/qwen3/export_dflash_draft.py
index 838ab14b6a5..75b72bd49cb 100644
--- a/examples/models/qwen3/export_dflash_draft.py
+++ b/examples/models/qwen3/export_dflash_draft.py
@@ -4,6 +4,7 @@
 import torch
 from huggingface_hub import snapshot_download
 from safetensors.torch import load_file
+from torch.export import Dim
 from transformers import AutoModelForCausalLM
 
 from executorch.backends.mlx.examples.llm.dflash_draft_model import DFlashDraftModel, load_dflash_config
@@ -36,21 +37,47 @@ def main():
     parser.add_argument("--output", default="qwen3_4b_dflash_draft.pte")
     parser.add_argument("--block-size", type=int, default=16)
     parser.add_argument("--ctx-len", type=int, default=8)
+    parser.add_argument("--max-ctx-len", type=int, default=4096)
     args = parser.parse_args()
 
     target = AutoModelForCausalLM.from_pretrained(args.target_model, dtype="auto")
     model = load_draft_model(args.draft_model, target.state_dict())
     model.eval()
-    model = model.float()
     del target
 
+    # Quantize to 4-bit to match the target export (--qlinear 4w --qembedding 4w,
+    # group_size=32). Without this the draft is full float32 (~3.5GB) with its
+    # shared embed/lm_head dominating memory; quantizing brings it to ~1GB and,
+    # critically, keeps the shared embed/lm_head at the SAME precision as the
+    # target so their logits stay consistent (acceptance depends on this).
+    from executorch.backends.mlx.llm.quantization import quantize_model_
+    quantize_model_(
+        model,
+        qlinear_config="4w",
+        qlinear_group_size=32,
+        qembedding_config="4w",
+        qembedding_group_size=32,
+        tie_word_embeddings=False,
+    )
+
     block_size, ctx_len = args.block_size, args.ctx_len
     hidden_size = model.fc.in_features
     tokens = torch.randint(0, 1000, (1, block_size), dtype=torch.long)
     target_hidden = torch.randn(1, ctx_len, hidden_size)
     position_ids = torch.arange(ctx_len + block_size).unsqueeze(0)
 
-    exported = torch.export.export(model, (tokens, target_hidden, position_ids))
+    ctx_dim = Dim("ctx_len", min=1, max=args.max_ctx_len)
+    dynamic_shapes = {
+        "tokens": None,
+        "target_hidden": {1: ctx_dim},
+        "position_ids": {1: ctx_dim + block_size},
+    }
+
+    import torch.fx.experimental._config as fx_config
+    with fx_config.patch(backed_size_oblivious=True):
+        exported = torch.export.export(
+            model, (tokens, target_hidden, position_ids), dynamic_shapes=dynamic_shapes
+        )
 
     from executorch.exir import to_edge_transform_and_lower
     from executorch.backends.mlx.partitioner import MLXPartitioner
@@ -61,6 +88,7 @@ def main():
     with open(args.output, "wb") as f:
         f.write(et_program.buffer)
     print(f"Saved draft model to: {args.output}")
+    print(f"Dynamic ctx_len supported: 1 to {args.max_ctx_len}, block_size fixed at {block_size}.")
 
 
 if __name__ == "__main__":
diff --git a/examples/models/qwen3/mlx_source_transformations.py b/examples/models/qwen3/mlx_source_transformations.py
new file mode 100644
index 00000000000..ee39af7bcda
--- /dev/null
+++ b/examples/models/qwen3/mlx_source_transformations.py
@@ -0,0 +1,110 @@
+"""Extracting Qwen3 hidden-state for DFlash.
+
+Same idea as examples/models/gemma4_31b/mlx_source_transformations.py --
+tap layers [2, N//2, N-3] and return them concatenated alongside logits.
+Gemma 4 does this by patching its own hand-written forward(). Qwen3 goes
+through the generic HF export path instead (export_llm_hf.py), which wraps
+the model in transformers' TorchExportableModuleWithStaticCache before
+torch.export. So we subclass that wrapper and add output_hidden_states
+to its forward rather than patching Qwen3 itself.
+
+Base class signature/behavior confirmed via:
+    inspect.getsource(transformers.integrations.executorch.TorchExportableModuleWithStaticCache)
+"""
+
+from typing import List, Optional, Sequence
+
+import torch
+from transformers.integrations.executorch import TorchExportableModuleWithStaticCache
+
+
+class TorchExportableModuleWithStaticCacheAndHidden(TorchExportableModuleWithStaticCache):
+    """forward() also returns tapped hidden states.
+    forward() -> (logits, hidden), where hidden is [B, T, len(layer_ids) * H].
+    """
+
+    def __init__(
+        self,
+        model,
+        batch_size: Optional[int] = None,
+        max_cache_len: Optional[int] = None,
+        device: Optional[torch.device] = None,
+        layer_ids: Sequence[int] = (),
+    ):
+        super().__init__(model, batch_size=batch_size, max_cache_len=max_cache_len, device=device)
+        if not layer_ids:
+            raise ValueError("layer_ids must be non-empty")
+        self.layer_ids: List[int] = list(layer_ids)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+    ):
+        outs = self.model(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            attention_mask=None,
+            past_key_values=self.static_cache,
+            use_cache=True,
+            output_hidden_states=True,
+        )
+
+        # hidden_states[0] is the embedding output, hidden_states[i+1] is decoder layer i's output
+        captured = [outs.hidden_states[i + 1] for i in self.layer_ids]
+        hidden = torch.cat(captured, dim=-1)
+
+        if hasattr(outs, "logits"):
+            return outs.logits, hidden
+        return outs.last_hidden_state, hidden
+
+
+def default_dflash_layer_ids(num_layers: int) -> List[int]:
+    """[2, N//2, N-3] tap pattern, same as Gemma 4. For Qwen3-4B (36 layers): [2, 18, 33]."""
+    return [2, num_layers // 2, num_layers - 3]
+
+class StatelessQwen3WithHidden(torch.nn.Module):
+    """Cache-free counterpart to TorchExportableModuleWithStaticCacheAndHidden.
+
+    DFlash's speculative-decode loop re-verifies overlapping/non-contiguous
+    token ranges every round (draft tokens get proposed, verified, some
+    rejected). TorchExportableModuleWithStaticCache's persistent internal
+    cache is built for strictly-sequential autoregressive decoding and
+    produces corrupted hidden states under this access pattern (confirmed at
+    the eager PyTorch level: two calls at non-contiguous cache_position values
+    on the same cached wrapper differ by ~8694 vs. ~0.0003 for a correctly
+    stateless model). This class sidesteps the whole problem: every forward()
+    call recomputes attention over exactly the tokens/positions given, with no
+    persistent state, matching the accumulate-full-context approach already
+    used by DFlashDraftModel in dflash_draft_model.py.
+
+    forward() -> (logits, hidden), where hidden is [B, T, len(layer_ids) * H].
+    """
+
+    def __init__(self, model, layer_ids: Sequence[int] = ()):
+        super().__init__()
+        if not layer_ids:
+            raise ValueError("layer_ids must be non-empty")
+        self.model = model
+        self.layer_ids: List[int] = list(layer_ids)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+    ):
+        outs = self.model(
+            input_ids=input_ids,
+            cache_position=cache_position,
+            attention_mask=None,
+            past_key_values=None,
+            use_cache=False,
+            output_hidden_states=True,
+        )
+        captured = [outs.hidden_states[i + 1] for i in self.layer_ids]
+        hidden = torch.cat(captured, dim=-1)
+        if hasattr(outs, "logits"):
+            return outs.logits, hidden
+        return outs.last_hidden_state, hidden
diff --git a/examples/models/qwen3/run_baseline.py b/examples/models/qwen3/run_baseline.py
new file mode 100644
index 00000000000..84a82ef8d5d
--- /dev/null
+++ b/examples/models/qwen3/run_baseline.py
@@ -0,0 +1,69 @@
+"""Plain autoregressive baseline using the SAME target .pte, tokenizer, and
+chat-template settings as run_dflash.py -- for an apples-to-apples comparison,
+not the old Phase 0 number (measured under different conditions/quant pass).
+"""
+import argparse
+import time
+
+import torch
+from transformers import AutoTokenizer
+from executorch.runtime import Runtime, Verification
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--target-pte", default="qwen3_4b_dflash_target.pte")
+    p.add_argument("--tokenizer", default="Qwen/Qwen3-4B")
+    p.add_argument("--prompt", required=True)
+    p.add_argument("--max-new-tokens", type=int, default=128)
+    p.add_argument("--chat-template", action="store_true", default=True)
+    p.add_argument("--no-chat-template", dest="chat_template", action="store_false")
+    p.add_argument("--enable-thinking", action="store_true", default=False)
+    args = p.parse_args()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, local_files_only=True)
+    eos_id = tokenizer.eos_token_id
+
+    if args.chat_template:
+        messages = [{"role": "user", "content": args.prompt}]
+        chat_out = tokenizer.apply_chat_template(
+            messages, add_generation_prompt=True,
+            enable_thinking=args.enable_thinking, return_tensors="pt",
+        )
+        prompt_ids = chat_out.input_ids if hasattr(chat_out, "input_ids") else chat_out
+    else:
+        prompt_ids = tokenizer(args.prompt, return_tensors="pt").input_ids
+
+    rt = Runtime.get()
+    target = rt.load_program(args.target_pte, verification=Verification.Minimal).load_method("forward")
+
+    prompt_len = prompt_ids.shape[1]
+    input_pos = torch.arange(prompt_len, dtype=torch.long)
+
+    t0 = time.time()
+    logits, _hidden = target.execute([prompt_ids, input_pos])
+    pos = prompt_len
+    token = int(logits[0, -1].argmax())
+    generated = [token]
+
+    while len(generated) < args.max_new_tokens:
+        tok_input = torch.tensor([[token]], dtype=torch.long)
+        pos_input = torch.tensor([pos], dtype=torch.long)
+        logits, _hidden = target.execute([tok_input, pos_input])
+        token = int(logits[0, -1].argmax())
+        generated.append(token)
+        pos += 1
+        if token == eos_id:
+            break
+
+    dt = time.time() - t0
+    text = tokenizer.decode(generated)
+    n = len(generated)
+    print(f"Prompt: {args.prompt}")
+    print(f"Generated ({n} tokens): {text}")
+    print(f"\n--- baseline stats ---")
+    print(f"time: {dt:.2f}s   tokens/s: {n / dt:.2f}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/models/qwen3/run_dflash.py b/examples/models/qwen3/run_dflash.py
new file mode 100644
index 00000000000..fd42c86aac5
--- /dev/null
+++ b/examples/models/qwen3/run_dflash.py
@@ -0,0 +1,162 @@
+"""DFlash speculative decoding driver for the ExecuTorch MLX backend (Python).
+
+Same four Phase 3 pieces as qwen3_dflash_engine.cpp, driven through the ET
+Python runtime so it runs on machines that can't build the C++ core:
+  1. draft block construction  [last_token, mask, mask, ...]
+  2. target verification        run target on [last_token] + draft_tokens
+  3. acceptance                 keep prefix up to first mismatch, + bonus token
+  4. position-based rollback     pos += accepted + 1
+
+V1 scope (per design doc): greedy, single batch, chain drafting, standard attn.
+"""
+
+import argparse
+import time
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from transformers import AutoTokenizer
+from executorch.runtime import Runtime, Verification
+
+from executorch.backends.mlx.examples.llm.dflash_draft_model import load_dflash_config
+
+
+def first_mismatch(draft_ids, target_ids):
+    """Number of leading draft tokens the target agrees with (greedy accept)."""
+    for i in range(len(draft_ids)):
+        if draft_ids[i] != target_ids[i]:
+            return i
+    return len(draft_ids)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--target-pte", default="qwen3_4b_dflash_target.pte")
+    p.add_argument("--draft-pte", default="qwen3_4b_dflash_draft.pte")
+    p.add_argument("--draft-model", default="z-lab/Qwen3-4B-DFlash-b16")
+    p.add_argument("--tokenizer", default="Qwen/Qwen3-4B")
+    p.add_argument("--prompt", default="The capital of France is")
+    p.add_argument("--max-new-tokens", type=int, default=64)
+    p.add_argument("--chat-template", action="store_true", default=True,
+                    help="Apply Qwen3's chat template (paper's eval setup). Default on.")
+    p.add_argument("--no-chat-template", dest="chat_template", action="store_false")
+    p.add_argument("--enable-thinking", action="store_true", default=False,
+                    help="Qwen3 thinking mode. Paper's Table 1 uses thinking mode DISABLED.")
+    p.add_argument("--block-size", type=int, default=None,
+                    help="Override the draft checkpoint config's block_size -- needed when "
+                         "--draft-pte was exported with a different block_size than the "
+                         "z-lab checkpoint's native config (e.g. our block_size=8 test export).")
+    args = p.parse_args()
+
+    config = load_dflash_config(Path(snapshot_download(
+        args.draft_model, allow_patterns=["*.json"], local_files_only=True)))
+    mask_id = config.mask_token_id
+    block_size = args.block_size if args.block_size is not None else config.block_size
+    layer_ids = config.target_layer_ids
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, local_files_only=True)
+    eos_id = tokenizer.eos_token_id
+
+    # Paper's Table 1 evaluates with the chat template applied and thinking mode
+    # disabled ("Q3-4B ... thinking mode disabled" -- Section 5.1), not raw
+    # completion text. The draft was also trained on prompt+response pairs
+    # (Section 5 "Datasets"; Figure 4 shows clean prompt p / response r), so
+    # feeding it untemplated text is a real distribution mismatch, not a
+    # cosmetic difference.
+
+    rt = Runtime.get()
+    target = rt.load_program(args.target_pte, verification=Verification.Minimal).load_method("forward")
+    draft = rt.load_program(args.draft_pte, verification=Verification.Minimal).load_method("forward")
+
+    if args.chat_template:
+        messages = [{"role": "user", "content": args.prompt}]
+        chat_out = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            enable_thinking=args.enable_thinking,
+            return_tensors="pt",
+        )
+        # Some transformers versions return a BatchEncoding (dict-like) here
+        # instead of a raw tensor; normalize either way.
+        prompt_ids = chat_out.input_ids if hasattr(chat_out, "input_ids") else chat_out
+    else:
+        prompt_ids = tokenizer(args.prompt, return_tensors="pt").input_ids
+    prompt_len = prompt_ids.shape[1]
+
+    # --- Prefill: target over the full prompt -> (logits, hidden) ---
+    input_pos = torch.arange(prompt_len, dtype=torch.long)
+    logits, hidden = target.execute([prompt_ids, input_pos])
+    hidden = hidden.float()
+    pos = prompt_len
+    last_token = int(logits[0, -1].argmax())
+
+    generated = [last_token]
+    rounds = 0
+    accepted_total = 0
+    t0 = time.time()
+
+    while len(generated) < args.max_new_tokens:
+        rounds += 1
+
+        # 1. Draft block: [last_token, mask, mask, ...]
+        draft_input = torch.cat(
+            [torch.tensor([[last_token]], dtype=torch.long),
+             torch.full((1, block_size - 1), mask_id, dtype=torch.long)], dim=1)
+        draft_pos = torch.arange(hidden.shape[1] + block_size, dtype=torch.long).unsqueeze(0)
+        _t0 = time.time()
+        (draft_logits,) = draft.execute([draft_input, hidden, draft_pos])
+        _draft_time = time.time() - _t0
+        draft_ids = draft_logits[0].argmax(-1).tolist()  # block_size - 1 tokens
+
+        # 2. Verify: target on [last_token] + draft_ids
+        verify_input = torch.cat(
+            [torch.tensor([[last_token]], dtype=torch.long),
+             torch.tensor([draft_ids], dtype=torch.long)], dim=1)
+        verify_pos = torch.arange(pos, pos + verify_input.shape[1], dtype=torch.long)
+        _t1 = time.time()
+        target_logits, new_hidden = target.execute([verify_input, verify_pos])
+        _verify_time = time.time() - _t1
+        if rounds <= 10:
+            print(f"  timing: draft={_draft_time*1000:.1f}ms  verify={_verify_time*1000:.1f}ms  ctx_len={hidden.shape[1]}")
+        target_ids = target_logits[0].argmax(-1).tolist()  # block_size tokens
+
+        # 3. Accept: matching prefix + the target's bonus token at the mismatch
+        accepted = first_mismatch(draft_ids, target_ids)
+        if rounds <= 5:
+            print(f"round {rounds}: pos={pos} hidden_ctx={hidden.shape[1]} "
+                  f"draft_ids[:5]={draft_ids[:5]} target_ids[:5]={target_ids[:5]} accepted={accepted}")
+        new_tokens = draft_ids[:accepted] + [target_ids[accepted]]
+        accepted_total += accepted
+
+        # Trim at EOS if it appears in the accepted run
+        if eos_id in new_tokens:
+            new_tokens = new_tokens[:new_tokens.index(eos_id) + 1]
+
+        generated.extend(new_tokens)
+
+        # 4. Position-based rollback
+        pos += len(new_tokens)
+        last_token = new_tokens[-1]
+        # Accumulate: append this round's newly-generated tokens' hidden onto
+        # the running context, don't replace it. The target context feature
+        # must span the whole sequence generated so far (paper Figure 2), not
+        # just the latest round.
+        hidden = torch.cat([hidden, new_hidden[:, :len(new_tokens), :].float()], dim=1)
+
+        if eos_id in new_tokens:
+            break
+
+    dt = time.time() - t0
+    text = tokenizer.decode(generated)
+    n = len(generated)
+    print(f"\nPrompt: {args.prompt}")
+    print(f"Generated ({n} tokens): {text}")
+    print(f"\n--- stats ---")
+    print(f"rounds: {rounds}")
+    print(f"avg accepted/round (tau proxy): {accepted_total / rounds:.2f}")
+    print(f"time: {dt:.2f}s   tokens/s: {n / dt:.2f}")
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/examples/models/qwen3/tests/test_dflash_draft.py b/examples/models/qwen3/tests/test_dflash_draft.py
new file mode 100644
index 00000000000..77dbd3d5fa6
--- /dev/null
+++ b/examples/models/qwen3/tests/test_dflash_draft.py
@@ -0,0 +1,28 @@
+"""Phase 2 / Part 8 verification: the exported draft .pte loads, executes, and
+supports dynamic ctx_len (since DFlash accumulated target-hidden context grows
+every speculative round -- see run_dflash.py). Covers export correctness and
+weight correctness implicitly: a shape or weight mismatch would have failed at
+export time (export_dflash_draft.py asserts checkpoint keys match before
+saving), so a successful load and execute here is sufficient proof.
+"""
+import sys
+import torch
+from executorch.runtime import Runtime, Verification
+
+pte_path = sys.argv[1] if len(sys.argv) > 1 else "qwen3_4b_dflash_draft.pte"
+et_runtime = Runtime.get()
+method = et_runtime.load_program(pte_path, verification=Verification.Minimal).load_method("forward")
+
+block_size, hidden_size, vocab_size = 16, 12800, 151936
+
+for ctx_len in (8, 20, 1):
+    tokens = torch.randint(0, 1000, (1, block_size), dtype=torch.long)
+    target_hidden = torch.randn(1, ctx_len, hidden_size)
+    position_ids = torch.arange(ctx_len + block_size).unsqueeze(0).long()
+
+    (draft_logits,) = method.execute([tokens, target_hidden, position_ids])
+    assert draft_logits.shape == (1, block_size - 1, vocab_size), (ctx_len, draft_logits.shape)
+    assert not torch.isnan(draft_logits).any() and not torch.isinf(draft_logits).any()
+    print(f"ctx_len={ctx_len}: OK {tuple(draft_logits.shape)}")
+
+print("PASS: draft .pte loads, executes, and supports dynamic ctx_len")
diff --git a/examples/models/qwen3/tests/test_dflash_draft_dynamic.py b/examples/models/qwen3/tests/test_dflash_draft_dynamic.py
new file mode 100644
index 00000000000..e1c5cf02b8a
--- /dev/null
+++ b/examples/models/qwen3/tests/test_dflash_draft_dynamic.py
@@ -0,0 +1,18 @@
+import torch
+from executorch.runtime import Runtime, Verification
+
+rt = Runtime.get()
+method = rt.load_program("qwen3_4b_dflash_draft.pte",
+                         verification=Verification.Minimal).load_method("forward")
+
+block_size, hidden_size, vocab_size = 16, 12800, 151936
+for ctx_len in (8, 20, 1):
+    tokens = torch.randint(0, 1000, (1, block_size), dtype=torch.long)
+    target_hidden = torch.randn(1, ctx_len, hidden_size)
+    position_ids = torch.arange(ctx_len + block_size).unsqueeze(0).long()
+    (draft_logits,) = method.execute([tokens, target_hidden, position_ids])
+    assert draft_logits.shape == (1, block_size - 1, vocab_size), (ctx_len, draft_logits.shape)
+    assert not torch.isnan(draft_logits).any() and not torch.isinf(draft_logits).any()
+    print(f"ctx_len={ctx_len}: OK {tuple(draft_logits.shape)}")
+
+print("PASS: dynamic ctx_len verified at multiple lengths")
diff --git a/examples/models/qwen3/tests/test_dflash_draft_eager.py b/examples/models/qwen3/tests/test_dflash_draft_eager.py
new file mode 100644
index 00000000000..ee3f0c06f8d
--- /dev/null
+++ b/examples/models/qwen3/tests/test_dflash_draft_eager.py
@@ -0,0 +1,34 @@
+import torch
+from executorch.backends.mlx.examples.llm.dflash_draft_model import DFlashConfig, DFlashDraftModel
+
+config = DFlashConfig(
+    hidden_size=2560,
+    num_hidden_layers=5,
+    num_attention_heads=32,
+    num_key_value_heads=8,
+    head_dim=128,
+    intermediate_size=9728,
+    vocab_size=151936,
+    rms_norm_eps=1e-6,
+    rope_theta=1_000_000.0,
+    max_position_embeddings=40960,
+    target_layer_ids=(1, 9, 17, 25, 33),
+    block_size=16,
+    mask_token_id=151669,
+    layer_types=("full_attention",) * 5,
+)
+
+model = DFlashDraftModel(config)
+model.eval()
+
+block_size, ctx_len = 16, 12
+tokens = torch.randint(0, config.vocab_size, (1, block_size), dtype=torch.long)
+target_hidden = torch.randn(1, ctx_len, len(config.target_layer_ids) * config.hidden_size)
+position_ids = torch.arange(ctx_len + block_size).unsqueeze(0)
+
+with torch.no_grad():
+    logits = model(tokens, target_hidden, position_ids)
+
+assert logits.shape == (1, block_size - 1, config.vocab_size), logits.shape
+assert not torch.isnan(logits).any() and not torch.isinf(logits).any()
+print(f"OK: draft logits {tuple(logits.shape)}, no NaN/Inf")
\ No newline at end of file
diff --git a/examples/models/qwen3/tests/test_dflash_draft_forward.py b/examples/models/qwen3/tests/test_dflash_draft_forward.py
new file mode 100644
index 00000000000..bfcc8895374
--- /dev/null
+++ b/examples/models/qwen3/tests/test_dflash_draft_forward.py
@@ -0,0 +1,46 @@
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from executorch.backends.mlx.examples.llm.dflash_draft_model import DFlashDraftModel, load_dflash_config
+
+path = Path(snapshot_download("z-lab/Qwen3-4B-DFlash-b16", allow_patterns=["*.safetensors", "*.json"]))
+config = load_dflash_config(path)
+
+model = DFlashDraftModel(config)
+weights = {}
+for f in path.glob("*.safetensors"):
+    weights.update(load_file(str(f)))
+model.load_state_dict(weights, strict=False)
+
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")
+target = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B", dtype="auto")
+model.embed_tokens.weight.data.copy_(target.model.embed_tokens.weight)
+model.lm_head.weight.data.copy_(target.lm_head.weight)
+model.eval()
+target.eval()
+
+prompt = "The capital of France is"
+input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+
+with torch.no_grad():
+    out = target(input_ids, output_hidden_states=True)
+    tapped = [out.hidden_states[i + 1] for i in config.target_layer_ids]
+    target_hidden = torch.cat(tapped, dim=-1).float()
+
+    last_token = input_ids[:, -1:]
+    block_size = 8
+    draft_tokens = torch.cat(
+        [last_token, torch.full((1, block_size - 1), config.mask_token_id, dtype=torch.long)], dim=1
+    )
+    position_ids = torch.arange(target_hidden.shape[1] + block_size).unsqueeze(0)
+
+    draft_logits = model(draft_tokens, target_hidden, position_ids)
+    predicted = draft_logits.argmax(-1)
+
+print("Prompt:", prompt)
+print("Predicted continuation:", tokenizer.decode(predicted[0]))
+print("Shape:", draft_logits.shape)
diff --git a/examples/models/qwen3/tests/test_dflash_draft_load_weights.py b/examples/models/qwen3/tests/test_dflash_draft_load_weights.py
new file mode 100644
index 00000000000..75af1d702db
--- /dev/null
+++ b/examples/models/qwen3/tests/test_dflash_draft_load_weights.py
@@ -0,0 +1,50 @@
+import json
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+from executorch.backends.mlx.examples.llm.dflash_draft_model import DFlashConfig, DFlashDraftModel
+
+path = Path(snapshot_download("z-lab/Qwen3-4B-DFlash-b16", allow_patterns=["*.safetensors", "*.json"]))
+cfg = json.loads((path / "config.json").read_text())
+dcfg = cfg["dflash_config"]
+
+config = DFlashConfig(
+    hidden_size=cfg["hidden_size"],
+    num_hidden_layers=cfg["num_hidden_layers"],
+    num_attention_heads=cfg["num_attention_heads"],
+    num_key_value_heads=cfg["num_key_value_heads"],
+    head_dim=cfg["head_dim"],
+    intermediate_size=cfg["intermediate_size"],
+    vocab_size=cfg["vocab_size"],
+    rms_norm_eps=cfg["rms_norm_eps"],
+    rope_theta=cfg["rope_theta"],
+    max_position_embeddings=cfg["max_position_embeddings"],
+    target_layer_ids=tuple(dcfg["target_layer_ids"]),
+    block_size=cfg["block_size"],
+    mask_token_id=dcfg["mask_token_id"],
+    layer_types=tuple(cfg.get("layer_types") or ["full_attention"] * cfg["num_hidden_layers"]),
+    sliding_window=cfg.get("sliding_window"),
+    final_logit_softcapping=cfg.get("final_logit_softcapping"),
+)
+
+model = DFlashDraftModel(config)
+
+draft_weights = {}
+for f in path.glob("*.safetensors"):
+    draft_weights.update(load_file(str(f)))
+
+print(f"Checkpoint has {len(draft_weights)} tensors")
+print("First 10 checkpoint keys:", list(draft_weights.keys())[:10])
+print("First 10 model keys:     ", list(model.state_dict().keys())[:10])
+
+missing, unexpected = model.load_state_dict(draft_weights, strict=False)
+still_missing = [k for k in missing if not k.startswith(("embed_tokens.", "lm_head."))]
+
+print(f"\nMissing (excl. embed/lm_head, expected empty): {still_missing}")
+print(f"Unexpected (expected empty): {unexpected}")
+
+assert not still_missing, "Architecture mismatch — key names don't match the real checkpoint"
+assert not unexpected, "Checkpoint has tensors our model doesn't define — architecture mismatch"
+print("\nOK: state_dict loaded cleanly, structure matches the real checkpoint")
\ No newline at end of file
diff --git a/examples/models/qwen3/tests/test_dflash_export.py b/examples/models/qwen3/tests/test_dflash_export.py
new file mode 100644
index 00000000000..17c6d7b0013
--- /dev/null
+++ b/examples/models/qwen3/tests/test_dflash_export.py
@@ -0,0 +1,32 @@
+"""Sanity check: qwen3_4b_dflash_target.pte returns (logits, hidden) with the
+expected shapes at runtime.
+
+Run after exporting with --dflash-layers, e.g.:
+    python3 export_llm_hf.py --model-id Qwen/Qwen3-4B --dflash-layers 2,18,33 ...
+    python3 test_dflash_export.py qwen3_4b_dflash_target.pte
+"""
+
+import sys
+import torch
+from executorch.runtime import Runtime, Verification
+
+DFLASH_LAYERS = [1, 9, 17, 25, 33]
+HIDDEN_SIZE = 2560
+EXPECTED_HIDDEN_DIM = len(DFLASH_LAYERS) * HIDDEN_SIZE  # 12800
+VOCAB_SIZE = 151936
+
+pte_path = sys.argv[1]
+et_runtime = Runtime.get()
+program = et_runtime.load_program(pte_path, verification=Verification.Minimal)
+method = program.load_method("forward")
+
+tokens = torch.tensor([[1, 2, 3]], dtype=torch.long)
+input_pos = torch.tensor([0], dtype=torch.long)
+logits, hidden = method.execute([tokens, input_pos])
+
+assert logits.shape == (1, 3, VOCAB_SIZE), logits.shape
+assert hidden.shape == (1, 3, EXPECTED_HIDDEN_DIM), hidden.shape
+assert not torch.isnan(logits).any() and not torch.isinf(logits).any()
+assert not torch.isnan(hidden).any() and not torch.isinf(hidden).any()
+
+print(f"OK: logits {tuple(logits.shape)}, hidden {tuple(hidden.shape)}")
\ No newline at end of file
diff --git a/examples/models/qwen3/tests/test_dflash_lossless.py b/examples/models/qwen3/tests/test_dflash_lossless.py
new file mode 100644
index 00000000000..86c81de5ece
--- /dev/null
+++ b/examples/models/qwen3/tests/test_dflash_lossless.py
@@ -0,0 +1,32 @@
+"""Verification (doc Part 8, criterion 1): greedy DFlash must be LOSSLESS --
+identical tokens to greedy autoregressive baseline, since the target verifies
+every accepted token. Any divergence means the speculative loop is broken."""
+import subprocess, sys, re
+
+PROMPT = "Write a Python function that takes a list of integers and returns the second largest number in the list."
+N = 96
+
+def run(script, extra):
+    out = subprocess.run(
+        [sys.executable, f"examples/models/qwen3/{script}",
+         "--prompt", PROMPT, "--max-new-tokens", str(N)] + extra,
+        capture_output=True, text=True, cwd=".",
+    ).stdout
+    m = re.search(r"Generated \([^)]*\): (.*?)\n\n", out, re.DOTALL)
+    return m.group(1) if m else out
+
+baseline = run("run_baseline.py", [])
+dflash = run("run_dflash.py", [])
+
+print("=== BASELINE ===\n", baseline[:400])
+print("\n=== DFLASH ===\n", dflash[:400])
+print("\n=== RESULT ===")
+if baseline.strip() == dflash.strip():
+    print("PASS: DFlash output is token-for-token identical to baseline (LOSSLESS)")
+else:
+    # find first divergence
+    for i, (a, b) in enumerate(zip(baseline, dflash)):
+        if a != b:
+            print(f"DIVERGE at char {i}: baseline={baseline[i:i+30]!r} dflash={dflash[i:i+30]!r}")
+            break
+    print("FAIL: outputs differ -- speculative loop is not lossless")
diff --git a/examples/models/qwen3/tests/test_dflash_target.py b/examples/models/qwen3/tests/test_dflash_target.py
new file mode 100644
index 00000000000..17c6d7b0013
--- /dev/null
+++ b/examples/models/qwen3/tests/test_dflash_target.py
@@ -0,0 +1,32 @@
+"""Sanity check: qwen3_4b_dflash_target.pte returns (logits, hidden) with the
+expected shapes at runtime.
+
+Run after exporting with --dflash-layers, e.g.:
+    python3 export_llm_hf.py --model-id Qwen/Qwen3-4B --dflash-layers 2,18,33 ...
+    python3 test_dflash_export.py qwen3_4b_dflash_target.pte
+"""
+
+import sys
+import torch
+from executorch.runtime import Runtime, Verification
+
+DFLASH_LAYERS = [1, 9, 17, 25, 33]
+HIDDEN_SIZE = 2560
+EXPECTED_HIDDEN_DIM = len(DFLASH_LAYERS) * HIDDEN_SIZE  # 12800
+VOCAB_SIZE = 151936
+
+pte_path = sys.argv[1]
+et_runtime = Runtime.get()
+program = et_runtime.load_program(pte_path, verification=Verification.Minimal)
+method = program.load_method("forward")
+
+tokens = torch.tensor([[1, 2, 3]], dtype=torch.long)
+input_pos = torch.tensor([0], dtype=torch.long)
+logits, hidden = method.execute([tokens, input_pos])
+
+assert logits.shape == (1, 3, VOCAB_SIZE), logits.shape
+assert hidden.shape == (1, 3, EXPECTED_HIDDEN_DIM), hidden.shape
+assert not torch.isnan(logits).any() and not torch.isinf(logits).any()
+assert not torch.isnan(hidden).any() and not torch.isinf(hidden).any()
+
+print(f"OK: logits {tuple(logits.shape)}, hidden {tuple(hidden.shape)}")
\ No newline at end of file

From 7be25cd611d7d43ac10934a53f7b3f4d0897d5be Mon Sep 17 00:00:00 2001
From: cthotti <cthotti@gmail.com>
Date: Fri, 10 Jul 2026 10:48:44 +0530
Subject: [PATCH 3/4] Add DFlash speculative decoding support for Qwen3

---
 .gitignore                                    |   7 +
 .../mlx/examples/llm/dflash_draft_model.py    | 118 ++++++++----
 backends/mlx/examples/llm/export_llm_hf.py    |  27 +--
 examples/models/qwen3/export_dflash_draft.py  |  38 ++--
 .../qwen3/mlx_source_transformations.py       |  64 ++-----
 examples/models/qwen3/run_baseline.py         |  19 +-
 examples/models/qwen3/run_dflash.py           | 177 ++++++++++++------
 .../models/qwen3/tests/test_dflash_draft.py   |  21 ++-
 .../qwen3/tests/test_dflash_draft_dynamic.py  |  18 --
 .../qwen3/tests/test_dflash_draft_eager.py    |  34 ----
 .../qwen3/tests/test_dflash_draft_forward.py  |  46 -----
 .../tests/test_dflash_draft_load_weights.py   |  50 -----
 .../qwen3/tests/test_dflash_draft_pte.py      |  19 --
 .../models/qwen3/tests/test_dflash_export.py  |  32 ----
 .../qwen3/tests/test_dflash_lossless.py       |  39 ++--
 .../models/qwen3/tests/test_dflash_target.py  |  11 +-
 16 files changed, 308 insertions(+), 412 deletions(-)
 delete mode 100644 examples/models/qwen3/tests/test_dflash_draft_dynamic.py
 delete mode 100644 examples/models/qwen3/tests/test_dflash_draft_eager.py
 delete mode 100644 examples/models/qwen3/tests/test_dflash_draft_forward.py
 delete mode 100644 examples/models/qwen3/tests/test_dflash_draft_load_weights.py
 delete mode 100644 examples/models/qwen3/tests/test_dflash_draft_pte.py
 delete mode 100644 examples/models/qwen3/tests/test_dflash_export.py

diff --git a/.gitignore b/.gitignore
index 10b048c17b5..69c6bcdc115 100644
--- a/.gitignore
+++ b/.gitignore
@@ -86,3 +86,10 @@ zephyr_dev_root.backup.*/
 # Agents
 .claude/*.local.*
 extension/pybindings/mlx.metallib
+
+# DFlash C++ engine (Qwen3) -- incomplete, broken build, follow-up PR
+examples/models/qwen3/CMakeLists.txt
+examples/models/qwen3/CMakePresets.json
+examples/models/qwen3/qwen3_dflash_engine.h
+examples/models/qwen3/qwen3_dflash_engine.cpp
+examples/models/qwen3/main.cpp
diff --git a/backends/mlx/examples/llm/dflash_draft_model.py b/backends/mlx/examples/llm/dflash_draft_model.py
index 913f89e2d57..225d330b9b9 100644
--- a/backends/mlx/examples/llm/dflash_draft_model.py
+++ b/backends/mlx/examples/llm/dflash_draft_model.py
@@ -1,26 +1,25 @@
-"""PyTorch DFlash draft model, structured for ExecuTorch export.
-
-Model-agnostic: the same code exports a valid draft .pte for any standard-
-attention target (Qwen3, Gemma-4, Llama-3.1) by reading a DFlashConfig loaded
-from the z-lab draft checkpoint. Per-model differences — RoPE base/scaling,
-sliding-window layers, final-logit softcap, embedding scale — are config values,
-so they resolve at trace time to one model-specific graph. The universal
-branches add no ops to the exported program.
-
-Two deliberate deviations from the reference forward, both for ET (design doc
-"Option A", self-contained draft):
-  - embed_tokens / lm_head are owned here (filled from the target at export)
-    instead of referenced live off the target module.
-  - forward returns draft logits (norm -> lm_head -> [:, 1:]) rather than the
-    bare normed hidden state; the reference does lm_head + logits_start=1 in its
-    generate loop.
-
-References:
-  z-lab/dflash  dflash/model_mlx.py  — universal MLX reference (Qwen3 / Qwen3.5 /
-      Gemma-4); source of the sliding-window, softcap, single-rope, QK-norm design.
-  z-lab/dflash  dflash/model.py       — PyTorch reference; exact weight layout for
-      load_state_dict.
-  transformers RoPE init               — config-driven inv_freq below.
+"""PyTorch implementation of the DFlash draft model for ExecuTorch export.
+
+This model is the lightweight "draft" network used in DFlash speculative
+decoding. Instead of generating one token at a time like the target LLM, it
+predicts an entire block of future tokens in parallel. To do this, it takes:
+    - proposal tokens (the draft block, beginning with the last accepted token),
+    - hidden states extracted from the target model (Phase 1), and
+    - position IDs for the draft block.
+
+The target hidden states are first projected into the draft model's hidden
+space, then every draft transformer layer attends to both the projected target
+context and the proposal tokens. The result is a fast approximation of what
+the target model is likely to generate next.
+
+The implementation is model-agnostic. Architectural differences such as RoPE,
+sliding-window attention, embedding scale, and logit softcapping come from
+DFlashConfig, allowing the same code to export draft models for Qwen3, Gemma,
+Llama, and other standard-attention architectures.
+
+For ExecuTorch export, the draft model owns its own embedding and LM head
+weights (copied from the target during export) and returns final draft logits
+directly rather than intermediate hidden states.
 """
 
 from dataclasses import dataclass, field
@@ -49,15 +48,18 @@ class DFlashConfig:
     layer_types: Tuple[str, ...] = field(default_factory=tuple)
     sliding_window: Optional[int] = None
     final_logit_softcapping: Optional[float] = None
-    embed_scale: float = 1.0  # 1.0 for Qwen3/Llama; sqrt(hidden_size) for Gemma
+    # Some models scale token embeddings before entering transformer.
+    # Qwen3/Llama use 1.0, while Gemma scales by sqrt(hidden_size).
+    embed_scale: float = 1.0
 
 
 def _rope_inv_freq(config: DFlashConfig) -> torch.Tensor:
-    # Covers the two rope types these drafts use: 'default' (Qwen3, Gemma-4) and
-    # 'linear' position-scaling. YaRN/longrope aren't handled — no current z-lab
-    # draft uses them.
+    # Build the RoPE frequencies expected by draft checkpoint.
+    # Different model families use different RoPE scaling strategies, so this is driven entirely from the checkpoint config rather than hardcoded.
     dim = config.head_dim
-    inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    inv_freq = 1.0 / (
+        config.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)
+    )
     scaling = config.rope_scaling or {}
     if scaling.get("rope_type", scaling.get("type")) == "linear":
         inv_freq = inv_freq / float(scaling["factor"])
@@ -72,7 +74,7 @@ def __init__(self, config: DFlashConfig):
     def forward(self, position_ids: torch.Tensor):
         inv = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
         pos = position_ids[:, None, :].float()
-        freqs = (inv @ pos).transpose(1, 2)      # [B, S, head_dim/2]
+        freqs = (inv @ pos).transpose(1, 2)  # [B, S, head_dim/2]
         emb = torch.cat((freqs, freqs), dim=-1)  # [B, S, head_dim]
         return emb.cos(), emb.sin()
 
@@ -81,15 +83,19 @@ def rotate_half(x: torch.Tensor) -> torch.Tensor:
     x1, x2 = x.chunk(2, dim=-1)
     return torch.cat((-x2, x1), dim=-1)
 
+
 def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    # Repeat KV heads when the model has fewer KV heads than attention heads.
     if n_rep == 1:
         return x
     b, h, s, d = x.shape
     x = x[:, :, None, :, :].expand(b, h, n_rep, s, d)
     return x.reshape(b, h * n_rep, s, d)
 
+
 def apply_rotary_pos_emb(q, k, cos, sin):
-    # q rotates over its own (last q_len) positions; k rotates over its full length.
+    # The proposal block only contains the current draft tokens, so queries rotate over those positions only.
+    # Keys contain both the target context and proposal block, so they rotate over the full combined sequence.
     q_len = q.shape[-2]
     cq, sq = cos[:, None, -q_len:, :], sin[:, None, -q_len:, :]
     ck, sk = cos[:, None, :, :], sin[:, None, :, :]
@@ -121,13 +127,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class DFlashAttention(nn.Module):
+    """Proposal tokens generate the queries. These represent the positions whose contents we are trying to predict."""
+
     def __init__(self, config: DFlashConfig, layer_idx: int):
         super().__init__()
         h, hd = config.hidden_size, config.head_dim
         self.n_heads = config.num_attention_heads
         self.n_kv = config.num_key_value_heads
         self.head_dim = hd
-        self.scaling = hd ** -0.5
+        self.scaling = hd**-0.5
         self.n_rep = self.n_heads // self.n_kv
         lt = config.layer_types
         self.is_sliding = bool(lt) and lt[layer_idx] == "sliding_attention"
@@ -140,11 +148,18 @@ def __init__(self, config: DFlashConfig, layer_idx: int):
         self.k_norm = DFlashRMSNorm(hd, config.rms_norm_eps)
 
     def forward(self, x, x_ctx, cos, sin):
+        """Keys and values come frmo both the projected target context and the proposal block itself. This lets the draft attend to what the target model already understands while also allowing predictions within the proposal block to interact with one another."""
         B, L, _ = x.shape
         S = x_ctx.shape[1]
-        q = self.q_norm(self.q_proj(x).view(B, L, self.n_heads, self.head_dim)).transpose(1, 2)
-        k = torch.cat([self.k_proj(x_ctx), self.k_proj(x)], dim=1).view(B, S + L, self.n_kv, self.head_dim)
-        v = torch.cat([self.v_proj(x_ctx), self.v_proj(x)], dim=1).view(B, S + L, self.n_kv, self.head_dim)
+        q = self.q_norm(
+            self.q_proj(x).view(B, L, self.n_heads, self.head_dim)
+        ).transpose(1, 2)
+        k = torch.cat([self.k_proj(x_ctx), self.k_proj(x)], dim=1).view(
+            B, S + L, self.n_kv, self.head_dim
+        )
+        v = torch.cat([self.v_proj(x_ctx), self.v_proj(x)], dim=1).view(
+            B, S + L, self.n_kv, self.head_dim
+        )
         k = self.k_norm(k).transpose(1, 2)
         v = v.transpose(1, 2)
         q, k = apply_rotary_pos_emb(q, k, cos, sin)
@@ -152,11 +167,14 @@ def forward(self, x, x_ctx, cos, sin):
             k = repeat_kv(k, self.n_rep)
             v = repeat_kv(v, self.n_rep)
         mask = self._sliding_mask(L, S, q.device, q.dtype) if self.is_sliding else None
+        # Each proposal position attends over the combined context to build a richer representation before predicting its token.
         out = torch.nn.functional.scaled_dot_product_attention(
-            q, k, v, attn_mask=mask, is_causal=False, scale=self.scaling)
+            q, k, v, attn_mask=mask, is_causal=False, scale=self.scaling
+        )
         return self.o_proj(out.transpose(1, 2).reshape(B, L, -1))
 
     def _sliding_mask(self, L, S, device, dtype):
+        """Restrict attention to the configured sliding window for models that use sliding-window attention, like Gemma."""
         total = S + L
         q_pos = torch.arange(S, total, device=device)[:, None]
         k_pos = torch.arange(total, device=device)[None, :]
@@ -170,9 +188,14 @@ def __init__(self, config: DFlashConfig, layer_idx: int):
         self.self_attn = DFlashAttention(config, layer_idx)
         self.mlp = DFlashMLP(config.hidden_size, config.intermediate_size)
         self.input_layernorm = DFlashRMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.post_attention_layernorm = DFlashRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = DFlashRMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
 
     def forward(self, x, x_ctx, cos, sin):
+        # Standard transformer decoder block:
+        # RMSNorm --> Attention --> Residual
+        # RMSNorm --> MLP --> Residual
         x = x + self.self_attn(self.input_layernorm(x), x_ctx, cos, sin)
         return x + self.mlp(self.post_attention_layernorm(x))
 
@@ -185,28 +208,39 @@ def __init__(self, config: DFlashConfig):
         self.fc = nn.Linear(concat_dim, config.hidden_size, bias=False)
         self.hidden_norm = DFlashRMSNorm(config.hidden_size, config.rms_norm_eps)
         self.layers = nn.ModuleList(
-            [DFlashDecoderLayer(config, i) for i in range(config.num_hidden_layers)])
+            [DFlashDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
         self.norm = DFlashRMSNorm(config.hidden_size, config.rms_norm_eps)
         self.rotary_emb = DFlashRotaryEmbedding(config)
-        # Option A: owned copies, filled from the target at export time.
+        # The draft owns its own embedding and LM head weights.
+        # During export these are copied from the target model, making the draft .pte self-contained.
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
     def forward(self, tokens, target_hidden, position_ids):
+        # Embed the proposal block (last accepted token + masked future positions).
         h = self.embed_tokens(tokens) * self.config.embed_scale
+        # Translate the concatenated target hidden states into the draft model's hidden space.
         h_ctx = self.hidden_norm(self.fc(target_hidden))
+        # Positional information for both the proposal block and target context.
         cos, sin = self.rotary_emb(position_ids)
         for layer in self.layers:
             h = layer(h, h_ctx, cos, sin)
         h = self.norm(h)
-        logits = self.lm_head(h[:, 1:, :])   # logits_start=1: drop the known first token
+        # Only return predictions for the future positions.
+        logits = self.lm_head(h[:, 1:, :])
+        # logits_start=1: drop the known first token
         cap = self.config.final_logit_softcapping
         if cap is not None:
             logits = torch.tanh(logits / cap) * cap
         return logits
 
+
 def load_dflash_config(checkpoint_dir) -> "DFlashConfig":
-    """Build a DFlashConfig from a z-lab DFlash checkpoint's config.json."""
+    """Load the architecture needed to reconstruct a DFlash draft model.
+
+    The checkpoint config describes both underlying transformer architecture (hidden size, attention heads, RoPE, etc.) and the DFlash-specific settings such as the tapped target layers and mask token.
+    """
     import json
     from pathlib import Path
 
@@ -227,7 +261,9 @@ def load_dflash_config(checkpoint_dir) -> "DFlashConfig":
         block_size=cfg["block_size"],
         mask_token_id=dcfg["mask_token_id"],
         rope_scaling=cfg.get("rope_scaling"),
-        layer_types=tuple(cfg.get("layer_types") or ["full_attention"] * cfg["num_hidden_layers"]),
+        layer_types=tuple(
+            cfg.get("layer_types") or ["full_attention"] * cfg["num_hidden_layers"]
+        ),
         sliding_window=cfg.get("sliding_window"),
         final_logit_softcapping=cfg.get("final_logit_softcapping"),
     )
diff --git a/backends/mlx/examples/llm/export_llm_hf.py b/backends/mlx/examples/llm/export_llm_hf.py
index 45fd8bfad28..91b081696c8 100644
--- a/backends/mlx/examples/llm/export_llm_hf.py
+++ b/backends/mlx/examples/llm/export_llm_hf.py
@@ -221,21 +221,18 @@ def _export_with_custom_components(
             max_cache_len=effective_cache_len,
         )
     elif dflash_layers is not None:
-        # Qwen3-specific for now
-        # Generalize the import if/when another model needs DFlash tapping.
-        #
-        # Stateless (no persistent KV cache), NOT TorchExportableModuleWithStaticCacheAndHidden:
-        # DFlash's speculative-decode loop re-verifies overlapping/non-contiguous token
-        # ranges every round, which corrupts a persistent StaticCache (confirmed at the
-        # eager PyTorch level -- see StatelessQwen3WithHidden's docstring). The driver
-        # always passes the full accumulated sequence instead of relying on a cache.
+        # Qwen3-specific for now.
         from executorch.examples.models.qwen3.mlx_source_transformations import (
-            StatelessQwen3WithHidden,
+            TorchExportableModuleWithStaticCacheAndHidden,
         )
 
-        logger.info(f"Creating stateless DFlash hidden-state-tapping wrapper, layers={dflash_layers}")
-        exportable = StatelessQwen3WithHidden(
+        logger.info(
+            f"Creating DFlash hidden-state-tapping wrapper, layers={dflash_layers}"
+        )
+        exportable = TorchExportableModuleWithStaticCacheAndHidden(
             model=model,
+            batch_size=1,
+            max_cache_len=effective_cache_len,
             layer_ids=dflash_layers,
         )
     else:
@@ -246,7 +243,7 @@ def _export_with_custom_components(
             max_cache_len=effective_cache_len,
         )
 
-    if use_custom_kv_cache and dflash_layers is None:
+    if use_custom_kv_cache:
         from executorch.backends.mlx.llm.source_transformation import (
             replace_hf_cache_with_mlx,
         )
@@ -318,6 +315,9 @@ def _export_with_custom_components(
         transform_passes=get_default_passes(),
         partitioner=[MLXPartitioner()],
         compile_config=edge_config,
+        # Required by the C++ LLMEngine metadata contract (get_llm_metadata in
+        # llm_runner_helper.cpp) -- this export path (used for --dflash-layers)
+        constant_methods={"get_max_seq_len": max_seq_len},
     )
 
     logger.info("Exporting to ExecuTorch...")
@@ -459,10 +459,11 @@ def main():
         "--dflash-layers",
         type=str,
         default=None,
-        help="Comma-separated layer indices to tap for DFlash hidden-state output, e.g. '2,18,33'",
+        help="Comma-separated transformer layer indices whose hidden states are concatenated and returned alongside logits for DFlash. E.g. '1,9,17,25,33'",
     )
 
     args = parser.parse_args()
+    # Convert "1,9,17,25,33" -> [1, 9, 17, 25, 33]
     dflash_layers = (
         [int(x) for x in args.dflash_layers.split(",")] if args.dflash_layers else None
     )
diff --git a/examples/models/qwen3/export_dflash_draft.py b/examples/models/qwen3/export_dflash_draft.py
index 75b72bd49cb..8b5fdc73291 100644
--- a/examples/models/qwen3/export_dflash_draft.py
+++ b/examples/models/qwen3/export_dflash_draft.py
@@ -1,14 +1,23 @@
+"""Exports the DFlash draft model to an .pte program. 
+
+This script loads the pretrained DFlash draft checkpoint, copies the shared embedding and output project weights from target model, applies same 4-bit quantization used by target, and exports the draft model for MLX inference. 
+The exported model is used alongside the target model during speculative decoding. 
+"""
+
 import argparse
 from pathlib import Path
 
 import torch
+
+from executorch.backends.mlx.examples.llm.dflash_draft_model import (
+    DFlashDraftModel,
+    load_dflash_config,
+)
 from huggingface_hub import snapshot_download
 from safetensors.torch import load_file
 from torch.export import Dim
 from transformers import AutoModelForCausalLM
 
-from executorch.backends.mlx.examples.llm.dflash_draft_model import DFlashDraftModel, load_dflash_config
-
 
 def load_draft_model(draft_id: str, target_state_dict: dict) -> DFlashDraftModel:
     path = Path(snapshot_download(draft_id, allow_patterns=["*.safetensors", "*.json"]))
@@ -21,11 +30,17 @@ def load_draft_model(draft_id: str, target_state_dict: dict) -> DFlashDraftModel
 
     missing, unexpected = model.load_state_dict(draft_weights, strict=False)
     assert not unexpected, f"Unexpected draft checkpoint keys: {unexpected}"
-    still_missing = [k for k in missing if not k.startswith(("embed_tokens.", "lm_head."))]
+    still_missing = [
+        k for k in missing if not k.startswith(("embed_tokens.", "lm_head."))
+    ]
     assert not still_missing, f"Missing draft checkpoint keys: {still_missing}"
 
     model.embed_tokens.weight.data.copy_(target_state_dict["model.embed_tokens.weight"])
-    lm_head_key = "lm_head.weight" if "lm_head.weight" in target_state_dict else "model.embed_tokens.weight"
+    lm_head_key = (
+        "lm_head.weight"
+        if "lm_head.weight" in target_state_dict
+        else "model.embed_tokens.weight"
+    )
     model.lm_head.weight.data.copy_(target_state_dict[lm_head_key])
     return model
 
@@ -45,12 +60,10 @@ def main():
     model.eval()
     del target
 
-    # Quantize to 4-bit to match the target export (--qlinear 4w --qembedding 4w,
-    # group_size=32). Without this the draft is full float32 (~3.5GB) with its
-    # shared embed/lm_head dominating memory; quantizing brings it to ~1GB and,
-    # critically, keeps the shared embed/lm_head at the SAME precision as the
-    # target so their logits stay consistent (acceptance depends on this).
+    # Quantize the draft model to match the target model.
+    # Keeping both models at the same precision reduces memory usage and helps keep their predictions consistent, which is important for achieving a high draft acceptance rate.
     from executorch.backends.mlx.llm.quantization import quantize_model_
+
     quantize_model_(
         model,
         qlinear_config="4w",
@@ -74,13 +87,14 @@ def main():
     }
 
     import torch.fx.experimental._config as fx_config
+
     with fx_config.patch(backed_size_oblivious=True):
         exported = torch.export.export(
             model, (tokens, target_hidden, position_ids), dynamic_shapes=dynamic_shapes
         )
 
-    from executorch.exir import to_edge_transform_and_lower
     from executorch.backends.mlx.partitioner import MLXPartitioner
+    from executorch.exir import to_edge_transform_and_lower
 
     edge = to_edge_transform_and_lower(exported, partitioner=[MLXPartitioner()])
     et_program = edge.to_executorch()
@@ -88,7 +102,9 @@ def main():
     with open(args.output, "wb") as f:
         f.write(et_program.buffer)
     print(f"Saved draft model to: {args.output}")
-    print(f"Dynamic ctx_len supported: 1 to {args.max_ctx_len}, block_size fixed at {block_size}.")
+    print(
+        f"Dynamic ctx_len supported: 1 to {args.max_ctx_len}, block_size fixed at {block_size}."
+    )
 
 
 if __name__ == "__main__":
diff --git a/examples/models/qwen3/mlx_source_transformations.py b/examples/models/qwen3/mlx_source_transformations.py
index ee39af7bcda..edfc41eff61 100644
--- a/examples/models/qwen3/mlx_source_transformations.py
+++ b/examples/models/qwen3/mlx_source_transformations.py
@@ -1,7 +1,9 @@
 """Extracting Qwen3 hidden-state for DFlash.
 
 Same idea as examples/models/gemma4_31b/mlx_source_transformations.py --
-tap layers [2, N//2, N-3] and return them concatenated alongside logits.
+extract layers and return them concatenated alongside logits. In Qwen3, the 
+layer ids from z-lab Qwen3 DFlash draft config is [1, 9, 17, 25, 33]
+
 Gemma 4 does this by patching its own hand-written forward(). Qwen3 goes
 through the generic HF export path instead (export_llm_hf.py), which wraps
 the model in transformers' TorchExportableModuleWithStaticCache before
@@ -18,10 +20,9 @@
 from transformers.integrations.executorch import TorchExportableModuleWithStaticCache
 
 
-class TorchExportableModuleWithStaticCacheAndHidden(TorchExportableModuleWithStaticCache):
-    """forward() also returns tapped hidden states.
-    forward() -> (logits, hidden), where hidden is [B, T, len(layer_ids) * H].
-    """
+class TorchExportableModuleWithStaticCacheAndHidden(
+    TorchExportableModuleWithStaticCache
+):
 
     def __init__(
         self,
@@ -31,7 +32,9 @@ def __init__(
         device: Optional[torch.device] = None,
         layer_ids: Sequence[int] = (),
     ):
-        super().__init__(model, batch_size=batch_size, max_cache_len=max_cache_len, device=device)
+        super().__init__(
+            model, batch_size=batch_size, max_cache_len=max_cache_len, device=device
+        )
         if not layer_ids:
             raise ValueError("layer_ids must be non-empty")
         self.layer_ids: List[int] = list(layer_ids)
@@ -52,7 +55,6 @@ def forward(
             output_hidden_states=True,
         )
 
-        # hidden_states[0] is the embedding output, hidden_states[i+1] is decoder layer i's output
         captured = [outs.hidden_states[i + 1] for i in self.layer_ids]
         hidden = torch.cat(captured, dim=-1)
 
@@ -62,49 +64,7 @@ def forward(
 
 
 def default_dflash_layer_ids(num_layers: int) -> List[int]:
-    """[2, N//2, N-3] tap pattern, same as Gemma 4. For Qwen3-4B (36 layers): [2, 18, 33]."""
-    return [2, num_layers // 2, num_layers - 3]
-
-class StatelessQwen3WithHidden(torch.nn.Module):
-    """Cache-free counterpart to TorchExportableModuleWithStaticCacheAndHidden.
-
-    DFlash's speculative-decode loop re-verifies overlapping/non-contiguous
-    token ranges every round (draft tokens get proposed, verified, some
-    rejected). TorchExportableModuleWithStaticCache's persistent internal
-    cache is built for strictly-sequential autoregressive decoding and
-    produces corrupted hidden states under this access pattern (confirmed at
-    the eager PyTorch level: two calls at non-contiguous cache_position values
-    on the same cached wrapper differ by ~8694 vs. ~0.0003 for a correctly
-    stateless model). This class sidesteps the whole problem: every forward()
-    call recomputes attention over exactly the tokens/positions given, with no
-    persistent state, matching the accumulate-full-context approach already
-    used by DFlashDraftModel in dflash_draft_model.py.
-
-    forward() -> (logits, hidden), where hidden is [B, T, len(layer_ids) * H].
+    """This is simply the default dflash layer selection for Qwen3.
+    [2, N//2, N-3] tap pattern, same as Gemma 4. For Qwen3-4B (36 layers): [2, 18, 33].
     """
-
-    def __init__(self, model, layer_ids: Sequence[int] = ()):
-        super().__init__()
-        if not layer_ids:
-            raise ValueError("layer_ids must be non-empty")
-        self.model = model
-        self.layer_ids: List[int] = list(layer_ids)
-
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        cache_position: Optional[torch.Tensor] = None,
-    ):
-        outs = self.model(
-            input_ids=input_ids,
-            cache_position=cache_position,
-            attention_mask=None,
-            past_key_values=None,
-            use_cache=False,
-            output_hidden_states=True,
-        )
-        captured = [outs.hidden_states[i + 1] for i in self.layer_ids]
-        hidden = torch.cat(captured, dim=-1)
-        if hasattr(outs, "logits"):
-            return outs.logits, hidden
-        return outs.last_hidden_state, hidden
+    return [2, num_layers // 2, num_layers - 3]
diff --git a/examples/models/qwen3/run_baseline.py b/examples/models/qwen3/run_baseline.py
index 84a82ef8d5d..c37b12b3b6d 100644
--- a/examples/models/qwen3/run_baseline.py
+++ b/examples/models/qwen3/run_baseline.py
@@ -1,13 +1,12 @@
-"""Plain autoregressive baseline using the SAME target .pte, tokenizer, and
-chat-template settings as run_dflash.py -- for an apples-to-apples comparison,
-not the old Phase 0 number (measured under different conditions/quant pass).
+"""Standard autoregressive decoding used as the baseline for the comparison. 
 """
+
 import argparse
 import time
 
 import torch
-from transformers import AutoTokenizer
 from executorch.runtime import Runtime, Verification
+from transformers import AutoTokenizer
 
 
 def main():
@@ -27,15 +26,19 @@ def main():
     if args.chat_template:
         messages = [{"role": "user", "content": args.prompt}]
         chat_out = tokenizer.apply_chat_template(
-            messages, add_generation_prompt=True,
-            enable_thinking=args.enable_thinking, return_tensors="pt",
+            messages,
+            add_generation_prompt=True,
+            enable_thinking=args.enable_thinking,
+            return_tensors="pt",
         )
         prompt_ids = chat_out.input_ids if hasattr(chat_out, "input_ids") else chat_out
     else:
         prompt_ids = tokenizer(args.prompt, return_tensors="pt").input_ids
 
     rt = Runtime.get()
-    target = rt.load_program(args.target_pte, verification=Verification.Minimal).load_method("forward")
+    target = rt.load_program(
+        args.target_pte, verification=Verification.Minimal
+    ).load_method("forward")
 
     prompt_len = prompt_ids.shape[1]
     input_pos = torch.arange(prompt_len, dtype=torch.long)
@@ -61,7 +64,7 @@ def main():
     n = len(generated)
     print(f"Prompt: {args.prompt}")
     print(f"Generated ({n} tokens): {text}")
-    print(f"\n--- baseline stats ---")
+    print("\n--baseline stats--")
     print(f"time: {dt:.2f}s   tokens/s: {n / dt:.2f}")
 
 
diff --git a/examples/models/qwen3/run_dflash.py b/examples/models/qwen3/run_dflash.py
index fd42c86aac5..849afcb874d 100644
--- a/examples/models/qwen3/run_dflash.py
+++ b/examples/models/qwen3/run_dflash.py
@@ -1,13 +1,18 @@
-"""DFlash speculative decoding driver for the ExecuTorch MLX backend (Python).
+"""Python implementation of the DFlash speculative decoding loop for the ExecuTorch MLX backend. 
 
-Same four Phase 3 pieces as qwen3_dflash_engine.cpp, driven through the ET
-Python runtime so it runs on machines that can't build the C++ core:
-  1. draft block construction  [last_token, mask, mask, ...]
-  2. target verification        run target on [last_token] + draft_tokens
-  3. acceptance                 keep prefix up to first mismatch, + bonus token
-  4. position-based rollback     pos += accepted + 1
+This file coordinates the interaction between the target model and the draft model during inference. Instead of asking the target model to generate one token at a time, DFlash first lets the lightweight draft model predict a block of future tokens, then asks the target model to verify those predictions in a single forward pass. Any matching draft tokens are accepted, while the first incorrect prediction is replaced with the target model's token. The process then repeats from the updated position. 
 
-V1 scope (per design doc): greedy, single batch, chain drafting, standard attn.
+Each speculation round consists of four steps: 
+    1. Build a draft block: [last_token, <mask>, <mask>, ...]
+    2. Run draft model to predict all masked tokens in parallel
+    3. Verify those predictions with the target model, keeping matching prefix and replacing the first mismatch with target's prediction. 
+    4. advance the sequence position to the newly accepted prefix and repeat. 
+
+V1 scope (per the issue discussion): 
+    - Greedy decoding
+    - Single-batch inference
+    - Chain drafting
+    - Standard attention models
 """
 
 import argparse
@@ -15,15 +20,15 @@
 from pathlib import Path
 
 import torch
-from huggingface_hub import snapshot_download
-from transformers import AutoTokenizer
-from executorch.runtime import Runtime, Verification
 
 from executorch.backends.mlx.examples.llm.dflash_draft_model import load_dflash_config
+from executorch.runtime import Runtime, Verification
+from huggingface_hub import snapshot_download
+from transformers import AutoTokenizer
 
 
 def first_mismatch(draft_ids, target_ids):
-    """Number of leading draft tokens the target agrees with (greedy accept)."""
+    """Returns the number of consecutive draft predictions that match the target."""
     for i in range(len(draft_ids)):
         if draft_ids[i] != target_ids[i]:
             return i
@@ -38,36 +43,57 @@ def main():
     p.add_argument("--tokenizer", default="Qwen/Qwen3-4B")
     p.add_argument("--prompt", default="The capital of France is")
     p.add_argument("--max-new-tokens", type=int, default=64)
-    p.add_argument("--chat-template", action="store_true", default=True,
-                    help="Apply Qwen3's chat template (paper's eval setup). Default on.")
+    p.add_argument(
+        "--chat-template",
+        action="store_true",
+        default=True,
+        help="Apply Qwen3's chat template (paper's eval setup). Default on.",
+    )
     p.add_argument("--no-chat-template", dest="chat_template", action="store_false")
-    p.add_argument("--enable-thinking", action="store_true", default=False,
-                    help="Qwen3 thinking mode. Paper's Table 1 uses thinking mode DISABLED.")
-    p.add_argument("--block-size", type=int, default=None,
-                    help="Override the draft checkpoint config's block_size -- needed when "
-                         "--draft-pte was exported with a different block_size than the "
-                         "z-lab checkpoint's native config (e.g. our block_size=8 test export).")
+    p.add_argument(
+        "--enable-thinking",
+        action="store_true",
+        default=False,
+        help="Qwen3 thinking mode. Paper's Table 1 uses thinking mode DISABLED.",
+    )
+    p.add_argument(
+        "--verbose",
+        action="store_true",
+        help="Print per-round timing/acceptance debug output.",
+    )
+    p.add_argument(
+        "--block-size",
+        type=int,
+        default=None,
+        help="Override the draft checkpoint config's block_size -- needed when "
+        "--draft-pte was exported with a different block_size than the "
+        "z-lab checkpoint's native config (e.g. our block_size=8 test export).",
+    )
     args = p.parse_args()
 
-    config = load_dflash_config(Path(snapshot_download(
-        args.draft_model, allow_patterns=["*.json"], local_files_only=True)))
+    config = load_dflash_config(
+        Path(
+            snapshot_download(
+                args.draft_model, allow_patterns=["*.json"], local_files_only=True
+            )
+        )
+    )
     mask_id = config.mask_token_id
     block_size = args.block_size if args.block_size is not None else config.block_size
-    layer_ids = config.target_layer_ids
 
     tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, local_files_only=True)
     eos_id = tokenizer.eos_token_id
 
-    # Paper's Table 1 evaluates with the chat template applied and thinking mode
-    # disabled ("Q3-4B ... thinking mode disabled" -- Section 5.1), not raw
-    # completion text. The draft was also trained on prompt+response pairs
-    # (Section 5 "Datasets"; Figure 4 shows clean prompt p / response r), so
-    # feeding it untemplated text is a real distribution mismatch, not a
-    # cosmetic difference.
+    # The draft model was trained on Qwen3 chat-formatted prompt/response pairs,so applying the same chat template during inference keeps the input distribution consistent with training.
+    # Using raw completion text noticeably reduces acceptance rates.
 
     rt = Runtime.get()
-    target = rt.load_program(args.target_pte, verification=Verification.Minimal).load_method("forward")
-    draft = rt.load_program(args.draft_pte, verification=Verification.Minimal).load_method("forward")
+    target = rt.load_program(
+        args.target_pte, verification=Verification.Minimal
+    ).load_method("forward")
+    draft = rt.load_program(
+        args.draft_pte, verification=Verification.Minimal
+    ).load_method("forward")
 
     if args.chat_template:
         messages = [{"role": "user", "content": args.prompt}]
@@ -77,14 +103,15 @@ def main():
             enable_thinking=args.enable_thinking,
             return_tensors="pt",
         )
-        # Some transformers versions return a BatchEncoding (dict-like) here
-        # instead of a raw tensor; normalize either way.
+        # Different Transformers versions return either a BatchEncoding or a tensor.
+        # Normalize both cases to a tensor.
         prompt_ids = chat_out.input_ids if hasattr(chat_out, "input_ids") else chat_out
     else:
         prompt_ids = tokenizer(args.prompt, return_tensors="pt").input_ids
     prompt_len = prompt_ids.shape[1]
 
-    # --- Prefill: target over the full prompt -> (logits, hidden) ---
+    # Run the target model over the prompt once to initialize generation.
+    # This produces the first next-token prediction and the hidden states that condition the draft model during speculative decoding.
     input_pos = torch.arange(prompt_len, dtype=torch.long)
     logits, hidden = target.execute([prompt_ids, input_pos])
     hidden = hidden.float()
@@ -94,55 +121,84 @@ def main():
     generated = [last_token]
     rounds = 0
     accepted_total = 0
+    emitted_total = 0
     t0 = time.time()
 
     while len(generated) < args.max_new_tokens:
         rounds += 1
+        # The exported draft model expects a fixed input shape for its token block.
+        # Although the hidden-state and position inputs support dynamic lengths, the token input does not.
+        # Reducing the block size near the end of generation causes the runtime to reject the input.
+        # Supporting truly dynamic block sizes would require exporting the draft model with a dynamic token dimension.
+        bs = block_size
 
-        # 1. Draft block: [last_token, mask, mask, ...]
+        # 1. Build draft input block.
         draft_input = torch.cat(
-            [torch.tensor([[last_token]], dtype=torch.long),
-             torch.full((1, block_size - 1), mask_id, dtype=torch.long)], dim=1)
-        draft_pos = torch.arange(hidden.shape[1] + block_size, dtype=torch.long).unsqueeze(0)
+            [
+                torch.tensor([[last_token]], dtype=torch.long),
+                torch.full((1, bs - 1), mask_id, dtype=torch.long),
+            ],
+            dim=1,
+        )
+        draft_pos = torch.arange(hidden.shape[1] + bs, dtype=torch.long).unsqueeze(0)
         _t0 = time.time()
         (draft_logits,) = draft.execute([draft_input, hidden, draft_pos])
-        _draft_time = time.time() - _t0
+        _draft_exec_time = time.time() - _t0
+        _t0b = time.time()
         draft_ids = draft_logits[0].argmax(-1).tolist()  # block_size - 1 tokens
+        _draft_argmax_time = time.time() - _t0b
 
-        # 2. Verify: target on [last_token] + draft_ids
+        # 2. Verify the draft predictions. Target model predicts the next token after every position in the block in a single forward pass.
         verify_input = torch.cat(
-            [torch.tensor([[last_token]], dtype=torch.long),
-             torch.tensor([draft_ids], dtype=torch.long)], dim=1)
+            [
+                torch.tensor([[last_token]], dtype=torch.long),
+                torch.tensor([draft_ids], dtype=torch.long),
+            ],
+            dim=1,
+        )
         verify_pos = torch.arange(pos, pos + verify_input.shape[1], dtype=torch.long)
         _t1 = time.time()
         target_logits, new_hidden = target.execute([verify_input, verify_pos])
-        _verify_time = time.time() - _t1
-        if rounds <= 10:
-            print(f"  timing: draft={_draft_time*1000:.1f}ms  verify={_verify_time*1000:.1f}ms  ctx_len={hidden.shape[1]}")
+        _target_exec_time = time.time() - _t1
+        _t1b = time.time()
         target_ids = target_logits[0].argmax(-1).tolist()  # block_size tokens
+        _target_argmax_time = time.time() - _t1b
 
-        # 3. Accept: matching prefix + the target's bonus token at the mismatch
+        # 3. Keep every drafting token that matches the target. At the first mismatch, stop accepting draft predictions and use the target model's token instead.
+        _t2 = time.time()
         accepted = first_mismatch(draft_ids, target_ids)
-        if rounds <= 5:
-            print(f"round {rounds}: pos={pos} hidden_ctx={hidden.shape[1]} "
-                  f"draft_ids[:5]={draft_ids[:5]} target_ids[:5]={target_ids[:5]} accepted={accepted}")
+        _fm_time = time.time() - _t2
+        if args.verbose and rounds <= 10:
+            print(
+                f"  timing: draft_exec={_draft_exec_time*1000:.1f}ms draft_argmax={_draft_argmax_time*1000:.2f}ms "
+                f"target_exec={_target_exec_time*1000:.1f}ms target_argmax={_target_argmax_time*1000:.2f}ms "
+                f"first_mismatch={_fm_time*1000:.3f}ms ctx_len={hidden.shape[1]}"
+            )
+        if args.verbose and rounds <= 5:
+            print(
+                f"round {rounds}: pos={pos} hidden_ctx={hidden.shape[1]} "
+                f"draft_ids[:5]={draft_ids[:5]} target_ids[:5]={target_ids[:5]} accepted={accepted}"
+            )
         new_tokens = draft_ids[:accepted] + [target_ids[accepted]]
         accepted_total += accepted
+        emitted_total += len(new_tokens)
 
-        # Trim at EOS if it appears in the accepted run
+        # Stop generation once an EOS token becomes part of the accepted sequence.
         if eos_id in new_tokens:
-            new_tokens = new_tokens[:new_tokens.index(eos_id) + 1]
+            new_tokens = new_tokens[: new_tokens.index(eos_id) + 1]
 
         generated.extend(new_tokens)
 
-        # 4. Position-based rollback
+        # 4. Advance the accepted sequence. Rejected draft tokens are discarded, and the next round starts from the updated position.
         pos += len(new_tokens)
         last_token = new_tokens[-1]
-        # Accumulate: append this round's newly-generated tokens' hidden onto
-        # the running context, don't replace it. The target context feature
-        # must span the whole sequence generated so far (paper Figure 2), not
-        # just the latest round.
-        hidden = torch.cat([hidden, new_hidden[:, :len(new_tokens), :].float()], dim=1)
+        # Append the hidden states for the newly accepted tokens to the running target context.
+        # The draft model conditions on the hidden states of the entire sequence, so this context grows as generation progresses rather than being replaced each round.
+        _t3 = time.time()
+        hidden = torch.cat([hidden, new_hidden[:, : len(new_tokens), :].float()], dim=1)
+        _cat_time = time.time() - _t3
+        if args.verbose and rounds <= 10:
+            print(f"  timing: hidden_cat={_cat_time*1000:.2f}ms")
 
         if eos_id in new_tokens:
             break
@@ -152,11 +208,12 @@ def main():
     n = len(generated)
     print(f"\nPrompt: {args.prompt}")
     print(f"Generated ({n} tokens): {text}")
-    print(f"\n--- stats ---")
+    print("\n--stats--")
     print(f"rounds: {rounds}")
-    print(f"avg accepted/round (tau proxy): {accepted_total / rounds:.2f}")
+    print(f"avg accepted/round (draft-only): {accepted_total / rounds:.2f}")
+    print(f"avg emitted/round (tau, incl. bonus): {emitted_total / rounds:.2f}")
     print(f"time: {dt:.2f}s   tokens/s: {n / dt:.2f}")
 
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
diff --git a/examples/models/qwen3/tests/test_dflash_draft.py b/examples/models/qwen3/tests/test_dflash_draft.py
index 77dbd3d5fa6..f01ea0d254e 100644
--- a/examples/models/qwen3/tests/test_dflash_draft.py
+++ b/examples/models/qwen3/tests/test_dflash_draft.py
@@ -1,17 +1,17 @@
-"""Phase 2 / Part 8 verification: the exported draft .pte loads, executes, and
-supports dynamic ctx_len (since DFlash accumulated target-hidden context grows
-every speculative round -- see run_dflash.py). Covers export correctness and
-weight correctness implicitly: a shape or weight mismatch would have failed at
-export time (export_dflash_draft.py asserts checkpoint keys match before
-saving), so a successful load and execute here is sufficient proof.
 """
+Verifies that the exported draft .pte loads, runs correctly, and supports dynamic context lengths as the accumulated target hidden-state context grows during speculative decoding. A successful export and execution also confirms that the checkpoint weights and exported model are compatible.
+"""
+
 import sys
+
 import torch
 from executorch.runtime import Runtime, Verification
 
 pte_path = sys.argv[1] if len(sys.argv) > 1 else "qwen3_4b_dflash_draft.pte"
 et_runtime = Runtime.get()
-method = et_runtime.load_program(pte_path, verification=Verification.Minimal).load_method("forward")
+method = et_runtime.load_program(
+    pte_path, verification=Verification.Minimal
+).load_method("forward")
 
 block_size, hidden_size, vocab_size = 16, 12800, 151936
 
@@ -21,8 +21,11 @@
     position_ids = torch.arange(ctx_len + block_size).unsqueeze(0).long()
 
     (draft_logits,) = method.execute([tokens, target_hidden, position_ids])
-    assert draft_logits.shape == (1, block_size - 1, vocab_size), (ctx_len, draft_logits.shape)
+    assert draft_logits.shape == (1, block_size - 1, vocab_size), (
+        ctx_len,
+        draft_logits.shape,
+    )
     assert not torch.isnan(draft_logits).any() and not torch.isinf(draft_logits).any()
     print(f"ctx_len={ctx_len}: OK {tuple(draft_logits.shape)}")
 
-print("PASS: draft .pte loads, executes, and supports dynamic ctx_len")
+print("PASS- draft .pte loads, executes, and supports dynamic ctx_len")
diff --git a/examples/models/qwen3/tests/test_dflash_draft_dynamic.py b/examples/models/qwen3/tests/test_dflash_draft_dynamic.py
deleted file mode 100644
index e1c5cf02b8a..00000000000
--- a/examples/models/qwen3/tests/test_dflash_draft_dynamic.py
+++ /dev/null
@@ -1,18 +0,0 @@
-import torch
-from executorch.runtime import Runtime, Verification
-
-rt = Runtime.get()
-method = rt.load_program("qwen3_4b_dflash_draft.pte",
-                         verification=Verification.Minimal).load_method("forward")
-
-block_size, hidden_size, vocab_size = 16, 12800, 151936
-for ctx_len in (8, 20, 1):
-    tokens = torch.randint(0, 1000, (1, block_size), dtype=torch.long)
-    target_hidden = torch.randn(1, ctx_len, hidden_size)
-    position_ids = torch.arange(ctx_len + block_size).unsqueeze(0).long()
-    (draft_logits,) = method.execute([tokens, target_hidden, position_ids])
-    assert draft_logits.shape == (1, block_size - 1, vocab_size), (ctx_len, draft_logits.shape)
-    assert not torch.isnan(draft_logits).any() and not torch.isinf(draft_logits).any()
-    print(f"ctx_len={ctx_len}: OK {tuple(draft_logits.shape)}")
-
-print("PASS: dynamic ctx_len verified at multiple lengths")
diff --git a/examples/models/qwen3/tests/test_dflash_draft_eager.py b/examples/models/qwen3/tests/test_dflash_draft_eager.py
deleted file mode 100644
index ee3f0c06f8d..00000000000
--- a/examples/models/qwen3/tests/test_dflash_draft_eager.py
+++ /dev/null
@@ -1,34 +0,0 @@
-import torch
-from executorch.backends.mlx.examples.llm.dflash_draft_model import DFlashConfig, DFlashDraftModel
-
-config = DFlashConfig(
-    hidden_size=2560,
-    num_hidden_layers=5,
-    num_attention_heads=32,
-    num_key_value_heads=8,
-    head_dim=128,
-    intermediate_size=9728,
-    vocab_size=151936,
-    rms_norm_eps=1e-6,
-    rope_theta=1_000_000.0,
-    max_position_embeddings=40960,
-    target_layer_ids=(1, 9, 17, 25, 33),
-    block_size=16,
-    mask_token_id=151669,
-    layer_types=("full_attention",) * 5,
-)
-
-model = DFlashDraftModel(config)
-model.eval()
-
-block_size, ctx_len = 16, 12
-tokens = torch.randint(0, config.vocab_size, (1, block_size), dtype=torch.long)
-target_hidden = torch.randn(1, ctx_len, len(config.target_layer_ids) * config.hidden_size)
-position_ids = torch.arange(ctx_len + block_size).unsqueeze(0)
-
-with torch.no_grad():
-    logits = model(tokens, target_hidden, position_ids)
-
-assert logits.shape == (1, block_size - 1, config.vocab_size), logits.shape
-assert not torch.isnan(logits).any() and not torch.isinf(logits).any()
-print(f"OK: draft logits {tuple(logits.shape)}, no NaN/Inf")
\ No newline at end of file
diff --git a/examples/models/qwen3/tests/test_dflash_draft_forward.py b/examples/models/qwen3/tests/test_dflash_draft_forward.py
deleted file mode 100644
index bfcc8895374..00000000000
--- a/examples/models/qwen3/tests/test_dflash_draft_forward.py
+++ /dev/null
@@ -1,46 +0,0 @@
-from pathlib import Path
-
-import torch
-from huggingface_hub import snapshot_download
-from safetensors.torch import load_file
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-from executorch.backends.mlx.examples.llm.dflash_draft_model import DFlashDraftModel, load_dflash_config
-
-path = Path(snapshot_download("z-lab/Qwen3-4B-DFlash-b16", allow_patterns=["*.safetensors", "*.json"]))
-config = load_dflash_config(path)
-
-model = DFlashDraftModel(config)
-weights = {}
-for f in path.glob("*.safetensors"):
-    weights.update(load_file(str(f)))
-model.load_state_dict(weights, strict=False)
-
-tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")
-target = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B", dtype="auto")
-model.embed_tokens.weight.data.copy_(target.model.embed_tokens.weight)
-model.lm_head.weight.data.copy_(target.lm_head.weight)
-model.eval()
-target.eval()
-
-prompt = "The capital of France is"
-input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-
-with torch.no_grad():
-    out = target(input_ids, output_hidden_states=True)
-    tapped = [out.hidden_states[i + 1] for i in config.target_layer_ids]
-    target_hidden = torch.cat(tapped, dim=-1).float()
-
-    last_token = input_ids[:, -1:]
-    block_size = 8
-    draft_tokens = torch.cat(
-        [last_token, torch.full((1, block_size - 1), config.mask_token_id, dtype=torch.long)], dim=1
-    )
-    position_ids = torch.arange(target_hidden.shape[1] + block_size).unsqueeze(0)
-
-    draft_logits = model(draft_tokens, target_hidden, position_ids)
-    predicted = draft_logits.argmax(-1)
-
-print("Prompt:", prompt)
-print("Predicted continuation:", tokenizer.decode(predicted[0]))
-print("Shape:", draft_logits.shape)
diff --git a/examples/models/qwen3/tests/test_dflash_draft_load_weights.py b/examples/models/qwen3/tests/test_dflash_draft_load_weights.py
deleted file mode 100644
index 75af1d702db..00000000000
--- a/examples/models/qwen3/tests/test_dflash_draft_load_weights.py
+++ /dev/null
@@ -1,50 +0,0 @@
-import json
-from pathlib import Path
-
-import torch
-from huggingface_hub import snapshot_download
-from safetensors.torch import load_file
-from executorch.backends.mlx.examples.llm.dflash_draft_model import DFlashConfig, DFlashDraftModel
-
-path = Path(snapshot_download("z-lab/Qwen3-4B-DFlash-b16", allow_patterns=["*.safetensors", "*.json"]))
-cfg = json.loads((path / "config.json").read_text())
-dcfg = cfg["dflash_config"]
-
-config = DFlashConfig(
-    hidden_size=cfg["hidden_size"],
-    num_hidden_layers=cfg["num_hidden_layers"],
-    num_attention_heads=cfg["num_attention_heads"],
-    num_key_value_heads=cfg["num_key_value_heads"],
-    head_dim=cfg["head_dim"],
-    intermediate_size=cfg["intermediate_size"],
-    vocab_size=cfg["vocab_size"],
-    rms_norm_eps=cfg["rms_norm_eps"],
-    rope_theta=cfg["rope_theta"],
-    max_position_embeddings=cfg["max_position_embeddings"],
-    target_layer_ids=tuple(dcfg["target_layer_ids"]),
-    block_size=cfg["block_size"],
-    mask_token_id=dcfg["mask_token_id"],
-    layer_types=tuple(cfg.get("layer_types") or ["full_attention"] * cfg["num_hidden_layers"]),
-    sliding_window=cfg.get("sliding_window"),
-    final_logit_softcapping=cfg.get("final_logit_softcapping"),
-)
-
-model = DFlashDraftModel(config)
-
-draft_weights = {}
-for f in path.glob("*.safetensors"):
-    draft_weights.update(load_file(str(f)))
-
-print(f"Checkpoint has {len(draft_weights)} tensors")
-print("First 10 checkpoint keys:", list(draft_weights.keys())[:10])
-print("First 10 model keys:     ", list(model.state_dict().keys())[:10])
-
-missing, unexpected = model.load_state_dict(draft_weights, strict=False)
-still_missing = [k for k in missing if not k.startswith(("embed_tokens.", "lm_head."))]
-
-print(f"\nMissing (excl. embed/lm_head, expected empty): {still_missing}")
-print(f"Unexpected (expected empty): {unexpected}")
-
-assert not still_missing, "Architecture mismatch — key names don't match the real checkpoint"
-assert not unexpected, "Checkpoint has tensors our model doesn't define — architecture mismatch"
-print("\nOK: state_dict loaded cleanly, structure matches the real checkpoint")
\ No newline at end of file
diff --git a/examples/models/qwen3/tests/test_dflash_draft_pte.py b/examples/models/qwen3/tests/test_dflash_draft_pte.py
deleted file mode 100644
index 374bb0e2aa0..00000000000
--- a/examples/models/qwen3/tests/test_dflash_draft_pte.py
+++ /dev/null
@@ -1,19 +0,0 @@
-import sys
-import torch
-from executorch.runtime import Runtime, Verification
-
-pte_path = sys.argv[1] if len(sys.argv) > 1 else "qwen3_4b_dflash_draft.pte"
-et_runtime = Runtime.get()
-program = et_runtime.load_program(pte_path, verification=Verification.Minimal)
-method = program.load_method("forward")
-
-# Must match the exact static shapes used at export time: ctx_len=8, block_size=16
-block_size, ctx_len, hidden_size, vocab_size = 16, 8, 12800, 151936
-tokens = torch.randint(0, 1000, (1, block_size), dtype=torch.long)
-target_hidden = torch.randn(1, ctx_len, hidden_size)
-position_ids = torch.arange(ctx_len + block_size).unsqueeze(0).long()
-
-(draft_logits,) = method.execute([tokens, target_hidden, position_ids])
-assert draft_logits.shape == (1, block_size - 1, vocab_size), draft_logits.shape
-assert not torch.isnan(draft_logits).any() and not torch.isinf(draft_logits).any()
-print(f"OK: draft_logits {tuple(draft_logits.shape)}")
diff --git a/examples/models/qwen3/tests/test_dflash_export.py b/examples/models/qwen3/tests/test_dflash_export.py
deleted file mode 100644
index 17c6d7b0013..00000000000
--- a/examples/models/qwen3/tests/test_dflash_export.py
+++ /dev/null
@@ -1,32 +0,0 @@
-"""Sanity check: qwen3_4b_dflash_target.pte returns (logits, hidden) with the
-expected shapes at runtime.
-
-Run after exporting with --dflash-layers, e.g.:
-    python3 export_llm_hf.py --model-id Qwen/Qwen3-4B --dflash-layers 2,18,33 ...
-    python3 test_dflash_export.py qwen3_4b_dflash_target.pte
-"""
-
-import sys
-import torch
-from executorch.runtime import Runtime, Verification
-
-DFLASH_LAYERS = [1, 9, 17, 25, 33]
-HIDDEN_SIZE = 2560
-EXPECTED_HIDDEN_DIM = len(DFLASH_LAYERS) * HIDDEN_SIZE  # 12800
-VOCAB_SIZE = 151936
-
-pte_path = sys.argv[1]
-et_runtime = Runtime.get()
-program = et_runtime.load_program(pte_path, verification=Verification.Minimal)
-method = program.load_method("forward")
-
-tokens = torch.tensor([[1, 2, 3]], dtype=torch.long)
-input_pos = torch.tensor([0], dtype=torch.long)
-logits, hidden = method.execute([tokens, input_pos])
-
-assert logits.shape == (1, 3, VOCAB_SIZE), logits.shape
-assert hidden.shape == (1, 3, EXPECTED_HIDDEN_DIM), hidden.shape
-assert not torch.isnan(logits).any() and not torch.isinf(logits).any()
-assert not torch.isnan(hidden).any() and not torch.isinf(hidden).any()
-
-print(f"OK: logits {tuple(logits.shape)}, hidden {tuple(hidden.shape)}")
\ No newline at end of file
diff --git a/examples/models/qwen3/tests/test_dflash_lossless.py b/examples/models/qwen3/tests/test_dflash_lossless.py
index 86c81de5ece..35d42f66084 100644
--- a/examples/models/qwen3/tests/test_dflash_lossless.py
+++ b/examples/models/qwen3/tests/test_dflash_lossless.py
@@ -1,32 +1,45 @@
-"""Verification (doc Part 8, criterion 1): greedy DFlash must be LOSSLESS --
-identical tokens to greedy autoregressive baseline, since the target verifies
-every accepted token. Any divergence means the speculative loop is broken."""
-import subprocess, sys, re
+"""Checks that DFlash prodcues the exact same output as normal greedy decoding."""
+
+import re
+import subprocess
+import sys
 
 PROMPT = "Write a Python function that takes a list of integers and returns the second largest number in the list."
 N = 96
 
+
 def run(script, extra):
     out = subprocess.run(
-        [sys.executable, f"examples/models/qwen3/{script}",
-         "--prompt", PROMPT, "--max-new-tokens", str(N)] + extra,
-        capture_output=True, text=True, cwd=".",
+        [
+            sys.executable,
+            f"examples/models/qwen3/{script}",
+            "--prompt",
+            PROMPT,
+            "--max-new-tokens",
+            str(N),
+        ]
+        + extra,
+        capture_output=True,
+        text=True,
+        cwd=".",
     ).stdout
     m = re.search(r"Generated \([^)]*\): (.*?)\n\n", out, re.DOTALL)
     return m.group(1) if m else out
 
+
 baseline = run("run_baseline.py", [])
 dflash = run("run_dflash.py", [])
 
-print("=== BASELINE ===\n", baseline[:400])
-print("\n=== DFLASH ===\n", dflash[:400])
-print("\n=== RESULT ===")
+print("BASELINE:\n", baseline[:400])
+print("\nDFLASH:\n", dflash[:400])
+print("\nRESULT:")
 if baseline.strip() == dflash.strip():
     print("PASS: DFlash output is token-for-token identical to baseline (LOSSLESS)")
 else:
-    # find first divergence
     for i, (a, b) in enumerate(zip(baseline, dflash)):
         if a != b:
-            print(f"DIVERGE at char {i}: baseline={baseline[i:i+30]!r} dflash={dflash[i:i+30]!r}")
+            print(
+                f"DIVERGE at char {i}: baseline={baseline[i:i+30]!r} dflash={dflash[i:i+30]!r}"
+            )
             break
-    print("FAIL: outputs differ -- speculative loop is not lossless")
+    print("FAIL- outputs differ: speculative loop is not lossless")
diff --git a/examples/models/qwen3/tests/test_dflash_target.py b/examples/models/qwen3/tests/test_dflash_target.py
index 17c6d7b0013..1cba7b803c9 100644
--- a/examples/models/qwen3/tests/test_dflash_target.py
+++ b/examples/models/qwen3/tests/test_dflash_target.py
@@ -1,12 +1,11 @@
-"""Sanity check: qwen3_4b_dflash_target.pte returns (logits, hidden) with the
-expected shapes at runtime.
+"""
+Verifies that the exported DFlash target model runs correctly and returns both logits and concatenated hidden states with the expected shapes. 
 
-Run after exporting with --dflash-layers, e.g.:
-    python3 export_llm_hf.py --model-id Qwen/Qwen3-4B --dflash-layers 2,18,33 ...
-    python3 test_dflash_export.py qwen3_4b_dflash_target.pte
+Run this after exporting the target model with --dflash-layers. 
 """
 
 import sys
+
 import torch
 from executorch.runtime import Runtime, Verification
 
@@ -29,4 +28,4 @@
 assert not torch.isnan(logits).any() and not torch.isinf(logits).any()
 assert not torch.isnan(hidden).any() and not torch.isinf(hidden).any()
 
-print(f"OK: logits {tuple(logits.shape)}, hidden {tuple(hidden.shape)}")
\ No newline at end of file
+print(f"OK- logits {tuple(logits.shape)}, hidden {tuple(hidden.shape)}")

From 01c4e8275d1cd8e576e4e38977b19d045a72ae9c Mon Sep 17 00:00:00 2001
From: Chet Hotti <cthotti@users.noreply.github.com>
Date: Tue, 14 Jul 2026 05:57:12 +0000
Subject: [PATCH 4/4] Address code review: drop dead Makefile target, fix
 pytest-collection hazard, trim profiling scaffolding, add license headers,
 fix typos

- Remove qwen3_dflash-mlx Makefile target + .gitignore block: depended on
  C++ engine files that are gitignored/not yet landed (follow-up PR)
- Rename tests/test_dflash_*.py -> check_dflash_*.py so pytest's test_*
  glob never collects these manual, hardware-gated driver scripts
- Trim sub-millisecond profiling scaffolding in run_dflash.py, keep only
  draft_exec/target_exec timing under --verbose (those dominate wall time)
- Add missing BSD license headers to 8 files
- Fix no-op --chat-template flag (only --no-chat-template had any effect)
- Fix typos: frmo->from, prodcues->produces, an .pte->a .pte,
  output project->output projection weights
- Add one-line comment on first_mismatch's draft/target length asymmetry
- Remove unused default_dflash_layer_ids (layer ids come from
  --dflash-layers / draft config, not this helper)
- Document DFlash's check_dflash_*.py as manual/CI-exempt in README
---
 .gitignore                                    |  8 ++---
 Makefile                                      | 12 +++----
 .../mlx/examples/llm/dflash_draft_model.py    |  8 ++++-
 examples/models/qwen3/README.md               | 22 +++++++++++++
 examples/models/qwen3/export_dflash_draft.py  | 10 ++++--
 .../qwen3/mlx_source_transformations.py       | 13 ++++----
 examples/models/qwen3/run_baseline.py         |  9 ++++--
 examples/models/qwen3/run_dflash.py           | 32 +++++++++----------
 ..._dflash_draft.py => check_dflash_draft.py} |  6 ++++
 ...h_lossless.py => check_dflash_lossless.py} |  8 ++++-
 ...flash_target.py => check_dflash_target.py} |  6 ++++
 11 files changed, 90 insertions(+), 44 deletions(-)
 rename examples/models/qwen3/tests/{test_dflash_draft.py => check_dflash_draft.py} (86%)
 rename examples/models/qwen3/tests/{test_dflash_lossless.py => check_dflash_lossless.py} (81%)
 rename examples/models/qwen3/tests/{test_dflash_target.py => check_dflash_target.py} (84%)

diff --git a/.gitignore b/.gitignore
index 69c6bcdc115..487b1d9e3ed 100644
--- a/.gitignore
+++ b/.gitignore
@@ -87,9 +87,5 @@ zephyr_dev_root.backup.*/
 .claude/*.local.*
 extension/pybindings/mlx.metallib
 
-# DFlash C++ engine (Qwen3) -- incomplete, broken build, follow-up PR
-examples/models/qwen3/CMakeLists.txt
-examples/models/qwen3/CMakePresets.json
-examples/models/qwen3/qwen3_dflash_engine.h
-examples/models/qwen3/qwen3_dflash_engine.cpp
-examples/models/qwen3/main.cpp
+# Scratch/WIP work not ready for review -- never committed
+/wip/
diff --git a/Makefile b/Makefile
index 26167aeb2b2..0891e06e1ba 100644
--- a/Makefile
+++ b/Makefile
@@ -485,11 +485,7 @@ clean:
 	       extension/llm/tokenizers/build \
 	       extension/llm/tokenizers/pytorch_tokenizers.egg-info
 
-qwen3_dflash-mlx:
-	@echo "==> Building and installing ExecuTorch with MLX..."
-	cmake --workflow --preset mlx-release
-	@echo "==> Building Qwen3 DFlash speculative decoding runner with MLX..."
-	cd examples/models/qwen3 && cmake --workflow --preset qwen3-dflash-mlx
-	@echo ""
-	@echo "✓ Build complete!"
-	@echo "  Runner: cmake-out/examples/models/qwen3/qwen3_dflash_runner"
+# qwen3_dflash-mlx target removed: it depended on the C++ engine sources
+# (CMakeLists.txt, CMakePresets.json, qwen3_dflash_engine.*), which are
+# gitignored/not yet landed. Restore this target in the follow-up PR that
+# actually lands the C++ engine.
diff --git a/backends/mlx/examples/llm/dflash_draft_model.py b/backends/mlx/examples/llm/dflash_draft_model.py
index 225d330b9b9..3562d840b72 100644
--- a/backends/mlx/examples/llm/dflash_draft_model.py
+++ b/backends/mlx/examples/llm/dflash_draft_model.py
@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 """PyTorch implementation of the DFlash draft model for ExecuTorch export.
 
 This model is the lightweight "draft" network used in DFlash speculative
@@ -148,7 +154,7 @@ def __init__(self, config: DFlashConfig, layer_idx: int):
         self.k_norm = DFlashRMSNorm(hd, config.rms_norm_eps)
 
     def forward(self, x, x_ctx, cos, sin):
-        """Keys and values come frmo both the projected target context and the proposal block itself. This lets the draft attend to what the target model already understands while also allowing predictions within the proposal block to interact with one another."""
+        """Keys and values come from both the projected target context and the proposal block itself. This lets the draft attend to what the target model already understands while also allowing predictions within the proposal block to interact with one another."""
         B, L, _ = x.shape
         S = x_ctx.shape[1]
         q = self.q_norm(
diff --git a/examples/models/qwen3/README.md b/examples/models/qwen3/README.md
index 123e65f16c5..2b3e3d94230 100644
--- a/examples/models/qwen3/README.md
+++ b/examples/models/qwen3/README.md
@@ -68,5 +68,27 @@ Note that you have to apply the chat template manually for the C++ runner.
 
 To run the model on an example iOS or Android app, see the Llama README's [Step 5: Build Mobile apps](../llama/README.md#step-5-build-mobile-apps) section.
 
+### DFlash speculative decoding (MLX delegate)
+
+`export_dflash_draft.py`, `run_dflash.py`, and `run_baseline.py` implement
+block-diffusion speculative decoding (DFlash) for Qwen3 on the MLX delegate.
+See `mlx_source_transformations.py` for the hidden-state-tapping wrapper used
+during export.
+
+The `check_dflash_*.py` scripts under `tests/` are manual driver scripts, not
+pytest tests -- they require exported `qwen3_4b_dflash_target.pte` /
+`_draft.pte` files (multi-GB, not checked in), HF downloads, and Apple
+M-series hardware with the MLX delegate, so they cannot run in this repo's
+CI. Run them by hand after exporting:
+
+```bash
+python examples/models/qwen3/tests/check_dflash_target.py qwen3_4b_dflash_target.pte
+python examples/models/qwen3/tests/check_dflash_draft.py qwen3_4b_dflash_draft.pte
+python examples/models/qwen3/tests/check_dflash_lossless.py
+```
+
+The "lossless" guarantee (DFlash output is token-for-token identical to
+greedy baseline decoding) is currently only verified this way, manually.
+
 ### FAQ
 For more help with exporting or running this model, feel free to ask in our [discord channel](https://discord.gg/UEjkY9Zs).
diff --git a/examples/models/qwen3/export_dflash_draft.py b/examples/models/qwen3/export_dflash_draft.py
index 8b5fdc73291..a3c18866764 100644
--- a/examples/models/qwen3/export_dflash_draft.py
+++ b/examples/models/qwen3/export_dflash_draft.py
@@ -1,6 +1,12 @@
-"""Exports the DFlash draft model to an .pte program. 
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
 
-This script loads the pretrained DFlash draft checkpoint, copies the shared embedding and output project weights from target model, applies same 4-bit quantization used by target, and exports the draft model for MLX inference. 
+"""Exports the DFlash draft model to a .pte program. 
+
+This script loads the pretrained DFlash draft checkpoint, copies the shared embedding and output projection weights from target model, applies same 4-bit quantization used by target, and exports the draft model for MLX inference. 
 The exported model is used alongside the target model during speculative decoding. 
 """
 
diff --git a/examples/models/qwen3/mlx_source_transformations.py b/examples/models/qwen3/mlx_source_transformations.py
index edfc41eff61..16a293c34e7 100644
--- a/examples/models/qwen3/mlx_source_transformations.py
+++ b/examples/models/qwen3/mlx_source_transformations.py
@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 """Extracting Qwen3 hidden-state for DFlash.
 
 Same idea as examples/models/gemma4_31b/mlx_source_transformations.py --
@@ -61,10 +67,3 @@ def forward(
         if hasattr(outs, "logits"):
             return outs.logits, hidden
         return outs.last_hidden_state, hidden
-
-
-def default_dflash_layer_ids(num_layers: int) -> List[int]:
-    """This is simply the default dflash layer selection for Qwen3.
-    [2, N//2, N-3] tap pattern, same as Gemma 4. For Qwen3-4B (36 layers): [2, 18, 33].
-    """
-    return [2, num_layers // 2, num_layers - 3]
diff --git a/examples/models/qwen3/run_baseline.py b/examples/models/qwen3/run_baseline.py
index c37b12b3b6d..594dada473a 100644
--- a/examples/models/qwen3/run_baseline.py
+++ b/examples/models/qwen3/run_baseline.py
@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 """Standard autoregressive decoding used as the baseline for the comparison. 
 """
 
@@ -15,8 +21,7 @@ def main():
     p.add_argument("--tokenizer", default="Qwen/Qwen3-4B")
     p.add_argument("--prompt", required=True)
     p.add_argument("--max-new-tokens", type=int, default=128)
-    p.add_argument("--chat-template", action="store_true", default=True)
-    p.add_argument("--no-chat-template", dest="chat_template", action="store_false")
+    p.add_argument("--no-chat-template", dest="chat_template", action="store_false", default=True)
     p.add_argument("--enable-thinking", action="store_true", default=False)
     args = p.parse_args()
 
diff --git a/examples/models/qwen3/run_dflash.py b/examples/models/qwen3/run_dflash.py
index 849afcb874d..d4c683ea86d 100644
--- a/examples/models/qwen3/run_dflash.py
+++ b/examples/models/qwen3/run_dflash.py
@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 """Python implementation of the DFlash speculative decoding loop for the ExecuTorch MLX backend. 
 
 This file coordinates the interaction between the target model and the draft model during inference. Instead of asking the target model to generate one token at a time, DFlash first lets the lightweight draft model predict a block of future tokens, then asks the target model to verify those predictions in a single forward pass. Any matching draft tokens are accepted, while the first incorrect prediction is replaced with the target model's token. The process then repeats from the updated position. 
@@ -44,12 +50,12 @@ def main():
     p.add_argument("--prompt", default="The capital of France is")
     p.add_argument("--max-new-tokens", type=int, default=64)
     p.add_argument(
-        "--chat-template",
-        action="store_true",
+        "--no-chat-template",
+        dest="chat_template",
+        action="store_false",
         default=True,
-        help="Apply Qwen3's chat template (paper's eval setup). Default on.",
+        help="Disable Qwen3's chat template. On by default (paper's eval setup).",
     )
-    p.add_argument("--no-chat-template", dest="chat_template", action="store_false")
     p.add_argument(
         "--enable-thinking",
         action="store_true",
@@ -144,9 +150,7 @@ def main():
         _t0 = time.time()
         (draft_logits,) = draft.execute([draft_input, hidden, draft_pos])
         _draft_exec_time = time.time() - _t0
-        _t0b = time.time()
         draft_ids = draft_logits[0].argmax(-1).tolist()  # block_size - 1 tokens
-        _draft_argmax_time = time.time() - _t0b
 
         # 2. Verify the draft predictions. Target model predicts the next token after every position in the block in a single forward pass.
         verify_input = torch.cat(
@@ -160,19 +164,17 @@ def main():
         _t1 = time.time()
         target_logits, new_hidden = target.execute([verify_input, verify_pos])
         _target_exec_time = time.time() - _t1
-        _t1b = time.time()
         target_ids = target_logits[0].argmax(-1).tolist()  # block_size tokens
-        _target_argmax_time = time.time() - _t1b
 
         # 3. Keep every drafting token that matches the target. At the first mismatch, stop accepting draft predictions and use the target model's token instead.
-        _t2 = time.time()
+        # (target_ids has block_size entries vs draft_ids' block_size - 1, so
+        # target_ids[accepted] is always in-bounds, including the all-accepted
+        # bonus-token case.)
         accepted = first_mismatch(draft_ids, target_ids)
-        _fm_time = time.time() - _t2
         if args.verbose and rounds <= 10:
             print(
-                f"  timing: draft_exec={_draft_exec_time*1000:.1f}ms draft_argmax={_draft_argmax_time*1000:.2f}ms "
-                f"target_exec={_target_exec_time*1000:.1f}ms target_argmax={_target_argmax_time*1000:.2f}ms "
-                f"first_mismatch={_fm_time*1000:.3f}ms ctx_len={hidden.shape[1]}"
+                f"  timing: draft_exec={_draft_exec_time*1000:.1f}ms "
+                f"target_exec={_target_exec_time*1000:.1f}ms ctx_len={hidden.shape[1]}"
             )
         if args.verbose and rounds <= 5:
             print(
@@ -194,11 +196,7 @@ def main():
         last_token = new_tokens[-1]
         # Append the hidden states for the newly accepted tokens to the running target context.
         # The draft model conditions on the hidden states of the entire sequence, so this context grows as generation progresses rather than being replaced each round.
-        _t3 = time.time()
         hidden = torch.cat([hidden, new_hidden[:, : len(new_tokens), :].float()], dim=1)
-        _cat_time = time.time() - _t3
-        if args.verbose and rounds <= 10:
-            print(f"  timing: hidden_cat={_cat_time*1000:.2f}ms")
 
         if eos_id in new_tokens:
             break
diff --git a/examples/models/qwen3/tests/test_dflash_draft.py b/examples/models/qwen3/tests/check_dflash_draft.py
similarity index 86%
rename from examples/models/qwen3/tests/test_dflash_draft.py
rename to examples/models/qwen3/tests/check_dflash_draft.py
index f01ea0d254e..a0c66b21986 100644
--- a/examples/models/qwen3/tests/test_dflash_draft.py
+++ b/examples/models/qwen3/tests/check_dflash_draft.py
@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 """
 Verifies that the exported draft .pte loads, runs correctly, and supports dynamic context lengths as the accumulated target hidden-state context grows during speculative decoding. A successful export and execution also confirms that the checkpoint weights and exported model are compatible.
 """
diff --git a/examples/models/qwen3/tests/test_dflash_lossless.py b/examples/models/qwen3/tests/check_dflash_lossless.py
similarity index 81%
rename from examples/models/qwen3/tests/test_dflash_lossless.py
rename to examples/models/qwen3/tests/check_dflash_lossless.py
index 35d42f66084..ea09a6bd196 100644
--- a/examples/models/qwen3/tests/test_dflash_lossless.py
+++ b/examples/models/qwen3/tests/check_dflash_lossless.py
@@ -1,4 +1,10 @@
-"""Checks that DFlash prodcues the exact same output as normal greedy decoding."""
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Checks that DFlash produces the exact same output as normal greedy decoding."""
 
 import re
 import subprocess
diff --git a/examples/models/qwen3/tests/test_dflash_target.py b/examples/models/qwen3/tests/check_dflash_target.py
similarity index 84%
rename from examples/models/qwen3/tests/test_dflash_target.py
rename to examples/models/qwen3/tests/check_dflash_target.py
index 1cba7b803c9..51f7e7e7f25 100644
--- a/examples/models/qwen3/tests/test_dflash_target.py
+++ b/examples/models/qwen3/tests/check_dflash_target.py
@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 """
 Verifies that the exported DFlash target model runs correctly and returns both logits and concatenated hidden states with the expected shapes.