diff --git a/.gitignore b/.gitignore
index ee206e23d94..487b1d9e3ed 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
@@ -85,3 +86,6 @@ zephyr_dev_root.backup.*/
 # Agents
 .claude/*.local.*
 extension/pybindings/mlx.metallib
+
+# Scratch/WIP work not ready for review -- never committed
+/wip/
diff --git a/Makefile b/Makefile
index 969b53644cd..0891e06e1ba 100644
--- a/Makefile
+++ b/Makefile
@@ -484,3 +484,8 @@ clean:
 	rm -rf cmake-out \
 	       extension/llm/tokenizers/build \
 	       extension/llm/tokenizers/pytorch_tokenizers.egg-info
+
+# 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/.gitignore b/backends/mlx/.gitignore
new file mode 100644
index 00000000000..697009f1b94
--- /dev/null
+++ b/backends/mlx/.gitignore
@@ -0,0 +1,9 @@
+# Auto-generated by backends/mlx/serialization/generate.py — do not commit.
+# See backends/mlx/serialization/README.md for regeneration instructions.
+runtime/MLXLoader.cpp
+runtime/MLXLoader.h
+runtime/schema_generated.h
+serialization/_generated/
+serialization/_generated_serializers.py
+serialization/mlx_graph_schema.py
+_generated_inspector.py
diff --git a/backends/mlx/examples/llm/dflash_draft_cache.py b/backends/mlx/examples/llm/dflash_draft_cache.py
new file mode 100644
index 00000000000..90aba011911
--- /dev/null
+++ b/backends/mlx/examples/llm/dflash_draft_cache.py
@@ -0,0 +1,70 @@
+# 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.
+
+"""Persistent per-layer KV cache for the DFlash draft model.
+
+Follows the TorchExportableModuleWithStaticCache pattern (per review): cache
+tensors are registered as mutable buffers and cache_position is passed
+through the call rather than tracked internally. Built on the existing
+KVCache (backends/mlx/llm/cache.py), so writes go through
+torch.ops.mlx.kv_cache_update instead of a Python slice (avoids the
+GuardOnDataDependentSymNode failure hit by an earlier attempt).
+"""
+
+from typing import Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from executorch.backends.mlx.llm.cache import KVCache
+
+
+class DFlashDraftKVCache(nn.Module):
+    def __init__(
+        self,
+        num_layers: int,
+        num_heads: int,
+        head_dim: int,
+        max_seq_len: int,
+        dtype: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+        self.max_seq_len = max_seq_len
+        self.layers = nn.ModuleList(
+            [
+                KVCache(
+                    max_batch_size=1,
+                    max_context_length=max_seq_len,
+                    n_heads=num_heads,
+                    head_dim=head_dim,
+                    enable_dynamic_shape=True,
+                    dtype=dtype,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def write(
+        self,
+        layer_idx: int,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        cache_position: Union[torch.Tensor, int],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Write K/V at cache_position and return the FULL buffer for this
+        # layer. Caller masks the unwritten tail via valid_mask(). Extract
+        # cache_position once and reuse across layers, not once per layer.
+        return self.layers[layer_idx].update(cache_position, key_states, value_states)
+
+    def valid_mask(self, valid_len: Union[torch.Tensor, int], device=None) -> torch.Tensor:
+        # True for positions [0, valid_len), False for the unwritten tail.
+        positions = torch.arange(self.max_seq_len, device=device)
+        return positions < valid_len
+
+    def reset(self) -> None:
+        for layer in self.layers:
+            layer.k_cache.zero_()
+            layer.v_cache.zero_()
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..bab038815e4
--- /dev/null
+++ b/backends/mlx/examples/llm/dflash_draft_model.py
@@ -0,0 +1,325 @@
+# 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
+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).
+
+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 (bidirectionally -- see DFlash paper Section
+4.2). The result is a fast approximation of what the target model is likely
+to generate next.
+
+Per review, this adapts HuggingFace's real Qwen3 building blocks (attention,
+RMSNorm, rotary embeddings, MLP, and HF's attention-interface dispatch --
+the same one the MLX integration registers "mlx" into, see
+backends/mlx/llm/hf_attention.py) rather than reimplementing them from
+scratch. Only the attention forward and decoder-layer container are
+DFlash-specific (queries come from the proposal block alone; keys/values
+span the projected target context concatenated with the block) -- everything
+else reuses the real HF modules directly, so behavior stays with the model
+implementation instead of drifting from a duplicated copy.
+
+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
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+from torch import nn
+
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.models.qwen3.modeling_qwen3 import (
+    apply_rotary_pos_emb,
+    eager_attention_forward,
+    Qwen3Attention,
+    Qwen3Config,
+    Qwen3MLP,
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+)
+
+
+@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
+    # 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 _to_qwen3_config(config: DFlashConfig) -> Qwen3Config:
+    """Translate DFlash's checkpoint-derived config into a real Qwen3Config,
+    so the draft model can build real Qwen3RMSNorm/Qwen3RotaryEmbedding/
+    Qwen3MLP/Qwen3Attention instances instead of hand-reimplemented copies.
+    """
+    rope_parameters = dict(config.rope_scaling or {})
+    rope_parameters["rope_type"] = rope_parameters.pop(
+        "type", rope_parameters.get("rope_type", "default")
+    )
+    rope_parameters["rope_theta"] = config.rope_theta
+
+    qwen3_config = Qwen3Config(
+        vocab_size=config.vocab_size,
+        hidden_size=config.hidden_size,
+        intermediate_size=config.intermediate_size,
+        num_hidden_layers=config.num_hidden_layers,
+        num_attention_heads=config.num_attention_heads,
+        num_key_value_heads=config.num_key_value_heads,
+        head_dim=config.head_dim,
+        max_position_embeddings=config.max_position_embeddings,
+        rms_norm_eps=config.rms_norm_eps,
+        rope_parameters=rope_parameters,
+        attention_bias=False,
+        layer_types=list(config.layer_types) or None,
+        sliding_window=config.sliding_window,
+    )
+    # Route through the same "mlx" attention interface the MLX integration
+    # registers for the target model (backends/mlx/llm/hf_attention.py),
+    # instead of calling torch SDPA directly.
+    qwen3_config._attn_implementation = "mlx"
+    return qwen3_config
+
+
+class DFlashQwen3Attention(Qwen3Attention):
+    """Adapts HF's Qwen3Attention to DFlash's cross-attention pattern:
+    queries come only from the proposal block, while keys/values span the
+    projected target context *and* the proposal block, concatenated. Reuses
+    the parent class's projections (q_proj/k_proj/v_proj/o_proj/q_norm/
+    k_norm), RoPE application (apply_rotary_pos_emb), and HF's attention
+    dispatch (ALL_ATTENTION_FUNCTIONS) rather than reimplementing any of
+    them.
+    """
+
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__(config, layer_idx)
+        # DFlash attends bidirectionally within a block and its target
+        # context (confirmed against the DFlash paper, Section 4.2: "Tokens
+        # attend bidirectionally within the same block and to the
+        # corresponding injected target context features"). This is never
+        # causal, unlike the base class's decode-time self-attention.
+        self.is_causal = False
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_ctx: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        cache=None,
+        cache_position=None,
+        new_ctx_len=None,
+    ) -> torch.Tensor:
+        # Two paths that MUST produce identical logits (verified in eager by
+        # scratch_draft_cache_equiv.py):
+        #   - uncached (cache=None): reproject the full [context; block]
+        #     every round. Original behavior, untouched.
+        #   - cached: reproject only the block plus the NEWLY-arrived context
+        #     tokens, write those context K/V into the per-layer cache at
+        #     cache_position, then read the full accumulated context K/V back
+        #     from the cache -- so context is projected once ever, not once
+        #     per round (removes the quadratic reprojection the review flagged).
+        B, L, _ = x.shape
+        S = x_ctx.shape[1]
+        q_shape = (B, L, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(x).view(q_shape)).transpose(1, 2)
+
+        total_len = cos.shape[1]
+        q_cos = cos.narrow(1, total_len - L, L)
+        q_sin = sin.narrow(1, total_len - L, L)
+        query_states, _ = apply_rotary_pos_emb(query_states, query_states, q_cos, q_sin)
+
+        if cache is None:
+            # --- Uncached path (unchanged) ---
+            kv_shape = (B, S + L, -1, self.head_dim)
+            kv_input = torch.cat([x_ctx, x], dim=1)
+            key_states = self.k_norm(self.k_proj(kv_input).view(kv_shape)).transpose(1, 2)
+            value_states = self.v_proj(kv_input).view(kv_shape).transpose(1, 2)
+            _, key_states = apply_rotary_pos_emb(key_states, key_states, cos, sin)
+        else:
+            # --- Cached path ---
+            # Project only newly-arrived context, RoPE at absolute positions,
+            # write to cache. A context token's position never changes, so
+            # caching post-RoPE keys is safe.
+            new_ctx = x_ctx.narrow(1, S - new_ctx_len, new_ctx_len)
+            new_shape = (B, new_ctx_len, -1, self.head_dim)
+            new_k = self.k_norm(self.k_proj(new_ctx).view(new_shape)).transpose(1, 2)
+            new_v = self.v_proj(new_ctx).view(new_shape).transpose(1, 2)
+            ctx_cos = cos.narrow(1, S - new_ctx_len, new_ctx_len)
+            ctx_sin = sin.narrow(1, S - new_ctx_len, new_ctx_len)
+            _, new_k = apply_rotary_pos_emb(new_k, new_k, ctx_cos, ctx_sin)
+
+            full_k, full_v = cache.write(self.layer_idx, new_k, new_v, cache_position)
+            ctx_key_states = full_k.narrow(2, 0, S)
+            ctx_value_states = full_v.narrow(2, 0, S)
+
+            # Block K/V are fresh every round (never cached).
+            blk_shape = (B, L, -1, self.head_dim)
+            blk_k = self.k_norm(self.k_proj(x).view(blk_shape)).transpose(1, 2)
+            blk_v = self.v_proj(x).view(blk_shape).transpose(1, 2)
+            blk_cos = cos.narrow(1, S, L)
+            blk_sin = sin.narrow(1, S, L)
+            _, blk_k = apply_rotary_pos_emb(blk_k, blk_k, blk_cos, blk_sin)
+
+            key_states = torch.cat([ctx_key_states, blk_k], dim=2)
+            value_states = torch.cat([ctx_value_states, blk_v], dim=2)
+
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, _ = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask=None,
+            dropout=0.0,
+            scaling=self.scaling,
+            sliding_window=None,
+        )
+        attn_output = attn_output.reshape(B, L, -1).contiguous()
+        return self.o_proj(attn_output)
+
+class DFlashQwen3DecoderLayer(nn.Module):
+    """Thin DFlash-specific container reusing real Qwen3 building blocks
+    (Qwen3RMSNorm, Qwen3MLP) plus the adapted DFlashQwen3Attention above. A
+    standard Qwen3DecoderLayer's forward() assumes one unified input
+    sequence; DFlash's split proposal-block/target-context query/key
+    pattern doesn't fit that contract, so this container still needs its
+    own forward(), but every actual computation is delegated to real HF
+    modules rather than reimplemented.
+    """
+
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.self_attn = DFlashQwen3Attention(config, layer_idx)
+        self.mlp = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(self, x, x_ctx, cos, sin, cache=None, cache_position=None, new_ctx_len=None):
+        x = x + self.self_attn(
+            self.input_layernorm(x), x_ctx, cos, sin,
+            cache=cache, cache_position=cache_position, new_ctx_len=new_ctx_len,
+        )
+        return x + self.mlp(self.post_attention_layernorm(x))
+
+class DFlashDraftModel(nn.Module):
+    def __init__(self, config: DFlashConfig):
+        super().__init__()
+        self.config = config
+        self.qwen3_config = _to_qwen3_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 = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.layers = nn.ModuleList(
+            [
+                DFlashQwen3DecoderLayer(self.qwen3_config, i)
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen3RotaryEmbedding(self.qwen3_config)
+        # 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, cache=None, cache_position=None, new_ctx_len=None):
+        # Positions are derived here from the actual input shapes rather than
+        # passed in as a separate tensor, so callers only need to supply
+        # `tokens` and `target_hidden` -- no third tensor whose shape must be
+        # kept in sync with both of theirs. This also keeps the block-length
+        # and context-length dimensions symbolically related for
+        # torch.export: since both are read directly off the dynamic-shaped
+        # inputs, the exporter ties them together automatically through this
+        # arithmetic instead of relying on a caller-supplied tensor with a
+        # separately-declared (and possibly mismatched) dynamic shape.
+        block_len = tokens.shape[1]
+        ctx_len = target_hidden.shape[1]
+        position_ids = torch.arange(ctx_len + block_len, device=tokens.device).unsqueeze(0)
+
+        # 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(h, position_ids)
+        for layer in self.layers:
+            h = layer(
+                h, h_ctx, cos, sin,
+                cache=cache, cache_position=cache_position, new_ctx_len=new_ctx_len,
+            )
+        h = self.norm(h)
+        # 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":
+    """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
+
+    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"),
+        embed_scale=cfg.get("embed_scale", dcfg.get("embed_scale", 1.0)),
+    )
diff --git a/backends/mlx/examples/llm/dflash_hidden_export.py b/backends/mlx/examples/llm/dflash_hidden_export.py
new file mode 100644
index 00000000000..ac9a30b5c83
--- /dev/null
+++ b/backends/mlx/examples/llm/dflash_hidden_export.py
@@ -0,0 +1,74 @@
+# 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.
+
+"""Generic hidden-state-tapping export wrapper for DFlash.
+
+Originally lived under examples/models/qwen3/, but the wrapper itself has
+no Qwen3-specific logic -- it subclasses transformers'
+TorchExportableModuleWithStaticCache and adds output_hidden_states to its
+forward, which works for any standard HF causal LM exported via the
+generic export_llm_hf.py path. Moved here (per review) so any model using
+that path can reuse it, rather than importing across from a
+model-specific folder.
+
+Gemma 4 (examples/models/gemma4_31b/) currently does hidden-state tapping
+differently -- by patching its own hand-written forward() rather than
+going through export_llm_hf.py's generic HF export path -- so it has its
+own separate mlx_source_transformations.py and isn't using this class.
+Not migrated as part of this change; that's a separate piece of work
+outside this PR's scope.
+
+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
+):
+
+    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,
+        )
+
+        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/backends/mlx/examples/llm/export_llm_hf.py b/backends/mlx/examples/llm/export_llm_hf.py
index fe6b8094f6b..a4ffecd03e1 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,21 @@ def _export_with_custom_components(
             batch_size=1,
             max_cache_len=effective_cache_len,
         )
+    elif dflash_layers is not None:
+        # Qwen3-specific for now.
+        from executorch.backends.mlx.examples.llm.dflash_hidden_export import (
+            TorchExportableModuleWithStaticCacheAndHidden,
+        )
+
+        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:
         logger.info("Creating TorchExportableModuleWithStaticCache wrapper...")
         exportable = TorchExportableModuleWithStaticCache(
@@ -299,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...")
@@ -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,18 @@ 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 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
+    )
 
     export_llama_hf(
         model_id=args.model_id,
@@ -450,6 +481,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/backends/mlx/llm/hf_attention.py b/backends/mlx/llm/hf_attention.py
index f2a01c9e653..cfb87b1115a 100644
--- a/backends/mlx/llm/hf_attention.py
+++ b/backends/mlx/llm/hf_attention.py
@@ -69,7 +69,13 @@ def mlx_sdpa_with_start_pos_forward(
         torch._check(start_pos + seq_len <= key.shape[2])
         attn_mask = None
     else:
-        start_pos = 0
+        # start_pos=0 would make stop_pos = start_pos + seq_len = seq_len,
+        # truncating key/value down to just the query's own length --
+        # silently discarding everything before it (e.g. DFlash's target
+        # context). Same fix pattern already used below in
+        # get_mlx_sliding_window_sdpa for the same underlying reason: set
+        # start_pos so stop_pos reaches the *full* key/value length instead.
+        start_pos = key.shape[2] - query.shape[2]
         attn_mask = attention_mask
 
     output = torch.ops.mlx.custom_sdpa(
diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py
index 002cda892f3..6ff9028b0f6 100644
--- a/backends/mlx/ops.py
+++ b/backends/mlx/ops.py
@@ -863,6 +863,18 @@ def handler(P: MLXProgramBuilder, n: Node) -> Slot:
     REGISTRY.register(target=[_target])(_make_scalar_int_handler(_node_cls, _op_name))
 
 
+@REGISTRY.register(target=[operator.neg])
+def _operator_neg_handler(P: MLXProgramBuilder, n: Node) -> Slot:
+    # operator.neg -> multiply by -1, reusing MultiplyIntNode.
+    args = P.args(n)
+    require_args(args, 1, 1, "operator.neg")
+    require_kwargs(P.kwargs(n), set(), "operator.neg")
+    (a,) = args
+    out = P.make_or_get_slot(n)
+    P.emit(MultiplyIntNode(a=P.to_int_or_vid(a), b=P.to_int_or_vid(-1), out=P.slot_to_vid(out)))
+    return out
+
+
 _REDUCTION_OPS: List[Tuple[List[Any], Any, str, int]] = [
     (
         [torch.ops.aten.sum.dim_IntList, torch.ops.aten.sum.default],
diff --git a/examples/models/qwen3/DFLASH_EXPERIMENTS.md b/examples/models/qwen3/DFLASH_EXPERIMENTS.md
new file mode 100644
index 00000000000..b1a5decd018
--- /dev/null
+++ b/examples/models/qwen3/DFLASH_EXPERIMENTS.md
@@ -0,0 +1,84 @@
+Written By: Chetan Thotti (cthotti)
+Date: 08/16/2026
+
+This is a record of the benchmarking we did on DFlash speculative decoding
+for Qwen3-4B, across three Apple Silicon machines. The short
+version: **DFlash's speedup depends on GPU architecture generation, not on
+how big or fast the chip otherwise is.** A base M4 clearly beats an M2 Pro
+here, even though the M2 Pro has more GPU cores and more memory bandwidth.
+If you're benchmarking DFlash on new hardware, read this first so you
+don't waste time on a chip that was never going to show a speedup.
+
+## The setup
+
+Model was `Qwen/Qwen3-4B` with the `z-lab/Qwen3-4B-DFlash-b16` draft
+checkpoint, exported with:
+--dflash-layers 1,9,17,25,33 --qlinear 4w --qembedding 4w --use-custom-sdpa --use-custom-kv-cache
+
+We tested three chips: the M2 in a MacBook Air (8 GPU cores), an M2 Pro
+rental (16 GPU cores), and a base M4 rental (10 GPU cores). Same `.pte`
+files got copied across the M2 Pro and M4 runs rather than re-exported
+separately, so any difference we saw was purely hardware, not export
+drift.
+
+## What we expected vs. what we found
+
+Going in, the assumption was that a "bigger" chip -- more GPU cores, more
+memory bandwidth -- would just be faster across the board, M2 Pro
+included. That's not what happened. Baseline (plain, one-token-at-a-time)
+decoding did scale the way you'd expect: M2 Pro's extra bandwidth made it
+faster than the M4 at baseline decoding, in every category we tested.
+But DFlash flipped that around entirely, only the M4 ever beat its own
+baseline. The M2 Pro was slower with DFlash turned on than without it,
+every single time.
+
+| Chip | Category | Baseline tok/s | DFlash tok/s | Speedup |
+|--------|------|-------|-------|-------|
+| M2 Air | Math | 25.65 | 19.44 | 0.76x |
+| M2 Air | Code | 28.19 | 23.81 | 0.84x |
+| M2 Air | Chat | 26.11 | 14.77 | 0.57x |
+| M2 Pro | Math | 48.30 | 42.63 | 0.88x |
+| M2 Pro | Code | 51.65 | 42.84 | 0.83x |
+| M2 Pro | Chat | 51.29 | 18.22 | 0.36x |
+| M4     | Math | 31.71 | 51.44 | 1.62x |
+| M4     | Code | 31.33 | 53.39 | 1.70x |
+| M4     | Chat | 31.42 | 22.73 | 0.72x |
+
+(Math/Code/Chat here are three different prompts, run 3 times each and
+averaged.)
+
+## The main difference between M2 and M4
+
+The performance difference comes down to GPU architecture, not the CPU.
+While I initially suspected SME2, the MLX backend doesn't use it. Instead,
+M3/M4 GPUs (Apple9) introduced Dynamic Caching and an improved SIMD matrix-
+multiply pipeline, which are much better suited for DFlash's verification
+stage. Dflash verifies an entire block of tokens at once using large matrix-
+matrix operations, allowing the M4 to execute this workload far more efficiently
+than the M2's Apple8 GPU, which lacks these architectural improvements. 
+
+**Practical takeaway: DFlash is worth using on M3 or M4 generation Macs,
+any tier, but not on an M1/M2.**
+
+## The draft model is also just bad at chat
+
+Separately from all the hardware stuff: math and code prompts got a tau
+(average tokens accepted per speculative round) around 6.6-6.8 on both
+chips. Chat-style prompts landed around 2.9 -- consistently, across three
+different chat prompts we tried, not just one unlucky example. Since tau
+was identical across both chips for the same prompt, this isn't a
+hardware issue at all, it's the draft model itself being noticeably
+worse at predicting open-ended conversational text than it is at
+structured math or code. Worth knowing if you're deciding whether DFlash
+is worth turning on for a particular kind of workload, independent of
+what hardware you're running it on.
+
+## What we didn't get to
+
+- **8-bit target quantization** (`--qlinear 8w`) fails to export with
+  `RuntimeError: Missing out variants: {'torchao::dequantize_affine'}`.
+  That's a real gap in the MLX partitioner's op coverage, not a flag
+  mistake.
+- **M3-generation chips** were never tested directly. Based on the
+  mechanism above they should behave like the M4 (same Apple9 GPU
+  family), but that's inference, not something we confirmed ourselves.
diff --git a/examples/models/qwen3/README.md b/examples/models/qwen3/README.md
index 123e65f16c5..0c410e39f6e 100644
--- a/examples/models/qwen3/README.md
+++ b/examples/models/qwen3/README.md
@@ -68,5 +68,28 @@ 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 `../../../backends/mlx/examples/llm/dflash_hidden_export.py` for the
+hidden-state-tapping wrapper used during export (moved out of this folder
+since it's model-agnostic, not Qwen3-specific).
+
+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
new file mode 100644
index 00000000000..e5af07c1571
--- /dev/null
+++ b/examples/models/qwen3/export_dflash_draft.py
@@ -0,0 +1,129 @@
+# 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.
+
+"""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. 
+"""
+
+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
+
+
+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():
+    # Register "mlx" into ALL_ATTENTION_FUNCTIONS so DFlashQwen3Attention's
+    # dispatch can resolve it. export_llm_hf.py does this for the target;
+    # the draft's own export script needs the same call.
+    from executorch.backends.mlx.llm.hf_attention import register_mlx_attention
+    register_mlx_attention()
+
+    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)
+    # --ctx-len only seeds the example shape used for tracing (ctx_len is a
+    # dynamic dim below via `Dim("ctx_len", ...)`); it is not a runtime cap.
+    # --max-ctx-len is the actual bound on context length at inference time.
+    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()
+    del target
+
+    # 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",
+        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)
+
+    # block_len is now bounded/dynamic; position_ids is no longer a model input.
+    ctx_dim = Dim("ctx_len", min=1, max=args.max_ctx_len)
+    # min=2, not 1: block_len=1 hits a shape-ambiguity guard during export.
+    # run_dflash.py's bs==1 branch already skips the draft entirely in that
+    # case (target-only step), so the export never needs to support it.
+    block_dim = Dim("block_len", min=2, max=block_size)
+    dynamic_shapes = {
+        "tokens": {1: block_dim},
+        "target_hidden": {1: ctx_dim},
+    }
+
+    import torch.fx.experimental._config as fx_config
+
+    with fx_config.patch(backed_size_oblivious=True):
+        exported = torch.export.export(
+            model, (tokens, target_hidden), dynamic_shapes=dynamic_shapes
+        )
+
+    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()
+
+    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}, dynamic block_len supported: 1 to {block_size}."
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/models/qwen3/run_baseline.py b/examples/models/qwen3/run_baseline.py
new file mode 100644
index 00000000000..f9df4e1bc1b
--- /dev/null
+++ b/examples/models/qwen3/run_baseline.py
@@ -0,0 +1,79 @@
+# 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. 
+"""
+
+import argparse
+import time
+
+import torch
+from executorch.runtime import Runtime, Verification
+from transformers import AutoTokenizer
+
+
+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(
+        "--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()
+
+    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("\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..8a69d2d98a5
--- /dev/null
+++ b/examples/models/qwen3/run_dflash.py
@@ -0,0 +1,237 @@
+# 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. 
+
+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
+import time
+from pathlib import Path
+
+import torch
+
+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):
+    """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
+    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(
+        "--no-chat-template",
+        dest="chat_template",
+        action="store_false",
+        default=True,
+        help="Disable Qwen3's chat template. On by default (paper's eval setup).",
+    )
+    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).",
+    )
+    p.add_argument("--max-seq-len", type=int, default=4096, help="Target's static cache capacity.")
+    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
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, local_files_only=True)
+    eos_id = tokenizer.eos_token_id
+
+    if tokenizer.chat_template is None:
+        # Some transformers versions expect chat templates in a standalone
+        # chat_template.jinja file and don't fall back to the legacy
+        # embedded tokenizer_config.json field when that file is absent.
+        import json
+        from huggingface_hub import hf_hub_download
+        cfg_path = hf_hub_download(args.tokenizer, "tokenizer_config.json")
+        cfg = json.loads(Path(cfg_path).read_text())
+        if "chat_template" in cfg:
+            tokenizer.chat_template = cfg["chat_template"]
+
+    # 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")
+
+    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",
+        )
+        # 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]
+
+    # 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()
+    pos = prompt_len
+    last_token = int(logits[0, -1].argmax())
+
+    generated = [last_token]
+    rounds = 0
+    accepted_total = 0
+    emitted_total = 0
+    t0 = time.time()
+
+    while len(generated) < args.max_new_tokens:
+        rounds += 1
+        # Dynamic block_len: shrink the proposal near the end of generation.
+        bs = min(block_size, args.max_new_tokens - len(generated), args.max_seq_len - pos)
+        if bs <= 0:
+            break
+
+        if bs == 1:
+            # No room to speculate -- skip the draft, target-only step.
+            draft_ids = []
+            _draft_exec_time = 0.0
+        else:
+            # 1. Build draft input block.
+            draft_input = torch.cat(
+                [
+                    torch.tensor([[last_token]], dtype=torch.long),
+                    torch.full((1, bs - 1), mask_id, dtype=torch.long),
+                ],
+                dim=1,
+            )
+            _t0 = time.time()
+            (draft_logits,) = draft.execute([draft_input, hidden])
+            _draft_exec_time = time.time() - _t0
+            draft_ids = draft_logits[0].argmax(-1).tolist()  # bs - 1 tokens
+
+        # 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,
+        )
+        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])
+        _target_exec_time = time.time() - _t1
+        target_ids = target_logits[0].argmax(-1).tolist()  # bs tokens
+
+        # 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.
+        # (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)
+        if args.verbose and rounds <= 10:
+            print(
+                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(
+                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]]
+
+        # Stop generation once an EOS token becomes part of the accepted sequence.
+        # Truncate before updating the running stats so a round that hits EOS
+        # doesn't over-count tokens/acceptances that never actually get emitted.
+        if eos_id in new_tokens:
+            new_tokens = new_tokens[: new_tokens.index(eos_id) + 1]
+            accepted = min(accepted, len(new_tokens) - 1)
+
+        accepted_total += accepted
+        emitted_total += len(new_tokens)
+
+        generated.extend(new_tokens)
+
+        # 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]
+        # 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.
+        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("\n--stats--")
+    print(f"rounds: {rounds}")
+    if rounds:
+        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()
diff --git a/examples/models/qwen3/tests/check_dflash_draft.py b/examples/models/qwen3/tests/check_dflash_draft.py
new file mode 100644
index 00000000000..a0c66b21986
--- /dev/null
+++ b/examples/models/qwen3/tests/check_dflash_draft.py
@@ -0,0 +1,37 @@
+# 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.
+"""
+
+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/check_dflash_draft_parity.py b/examples/models/qwen3/tests/check_dflash_draft_parity.py
new file mode 100644
index 00000000000..371390327f0
--- /dev/null
+++ b/examples/models/qwen3/tests/check_dflash_draft_parity.py
@@ -0,0 +1,81 @@
+# 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.
+
+"""Eager parity test: DFlashQwen3Attention against real, unmodified
+Qwen3Attention. Run directly:
+    python examples/models/qwen3/tests/check_dflash_draft_parity.py
+
+Property under test: with an empty target context (S=0), DFlash's attention
+reduces exactly to plain bidirectional self-attention over the block alone --
+what a real Qwen3Attention computes with is_causal=False. Divergence here
+means the adaptation (position slicing, masking, dispatch) has drifted from
+the reference module.
+"""
+
+import copy
+
+import torch
+
+from executorch.backends.mlx.examples.llm.dflash_draft_model import (
+    DFlashConfig,
+    DFlashQwen3Attention,
+    _to_qwen3_config,
+)
+from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention, Qwen3RotaryEmbedding
+
+
+def main():
+    torch.manual_seed(0)
+
+    config = DFlashConfig(
+        hidden_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        intermediate_size=128,
+        vocab_size=100,
+        rms_norm_eps=1e-6,
+        rope_theta=10000.0,
+        max_position_embeddings=128,
+        target_layer_ids=(0,),
+    )
+    qwen3_config = _to_qwen3_config(config)
+    qwen3_config._attn_implementation = "eager"
+
+    dflash_attn = DFlashQwen3Attention(qwen3_config, layer_idx=0).eval()
+    ref_attn = Qwen3Attention(qwen3_config, layer_idx=0).eval()
+    ref_attn.is_causal = False
+
+    ref_attn.load_state_dict(copy.deepcopy(dflash_attn.state_dict()))
+
+    B, L = 1, 5
+    x = torch.randn(B, L, config.hidden_size)
+    x_ctx = torch.zeros(B, 0, config.hidden_size)
+
+    rotary_emb = Qwen3RotaryEmbedding(qwen3_config)
+    position_ids = torch.arange(L).unsqueeze(0)
+    cos, sin = rotary_emb(x, position_ids)
+
+    dflash_out = dflash_attn(x, x_ctx, cos, sin)
+    ref_out, _ = ref_attn(
+        hidden_states=x,
+        position_embeddings=(cos, sin),
+        attention_mask=None,
+        past_key_values=None,
+        cache_position=None,
+    )
+
+    max_diff = (dflash_out - ref_out).abs().max().item()
+    print(f"max abs diff: {max_diff:.3e}")
+    if torch.allclose(dflash_out, ref_out, atol=1e-5, rtol=1e-5):
+        print("PASS: DFlashQwen3Attention matches real Qwen3Attention with empty context.")
+    else:
+        print("FAIL: DFlashQwen3Attention diverges from the reference module.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/models/qwen3/tests/check_dflash_lossless.py b/examples/models/qwen3/tests/check_dflash_lossless.py
new file mode 100644
index 00000000000..ea09a6bd196
--- /dev/null
+++ b/examples/models/qwen3/tests/check_dflash_lossless.py
@@ -0,0 +1,51 @@
+# 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
+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=".",
+    ).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("\nDFLASH:\n", dflash[:400])
+print("\nRESULT:")
+if baseline.strip() == dflash.strip():
+    print("PASS: DFlash output is token-for-token identical to baseline (LOSSLESS)")
+else:
+    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/check_dflash_target.py b/examples/models/qwen3/tests/check_dflash_target.py
new file mode 100644
index 00000000000..51f7e7e7f25
--- /dev/null
+++ b/examples/models/qwen3/tests/check_dflash_target.py
@@ -0,0 +1,37 @@
+# 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. 
+
+Run this after exporting the target model with --dflash-layers. 
+"""
+
+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)}")
diff --git a/scratch_draft_cache_equiv.py b/scratch_draft_cache_equiv.py
new file mode 100644
index 00000000000..adfe80dd5d1
--- /dev/null
+++ b/scratch_draft_cache_equiv.py
@@ -0,0 +1,111 @@
+"""Eager equivalence probe for the DFlash draft KV cache.
+
+Not part of the PR -- a throwaway diagnostic. Run directly:
+    python scratch_draft_cache_equiv.py
+
+Proves the cached draft forward produces logits numerically identical to
+the uncached path across several simulated speculative-decoding rounds.
+"""
+
+import torch
+
+from executorch.backends.mlx.examples.llm.dflash_draft_model import (
+    DFlashConfig,
+    DFlashDraftModel,
+)
+
+
+def make_model():
+    torch.manual_seed(0)
+    config = DFlashConfig(
+        hidden_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        intermediate_size=128,
+        vocab_size=100,
+        rms_norm_eps=1e-6,
+        rope_theta=10000.0,
+        max_position_embeddings=128,
+        target_layer_ids=(0,),
+        block_size=8,
+    )
+    model = DFlashDraftModel(config)
+    model.qwen3_config._attn_implementation = "eager"
+    for layer in model.layers:
+        layer.self_attn.config._attn_implementation = "eager"
+    return model.eval(), config
+
+
+def simulate_rounds(model, config, num_rounds=4, use_cache=False):
+    hidden_size = config.hidden_size
+    block_len = config.block_size
+
+    gen = torch.Generator().manual_seed(123)
+    ctx = torch.randn(1, 5, hidden_size, generator=gen)
+    per_round_logits = []
+
+    cache = None
+    if use_cache:
+        from executorch.backends.mlx.examples.llm.dflash_draft_cache import (
+            DFlashDraftKVCache,
+        )
+
+        cache = DFlashDraftKVCache(
+            num_layers=config.num_hidden_layers,
+            num_heads=config.num_key_value_heads,
+            head_dim=config.head_dim,
+            max_seq_len=config.max_position_embeddings,
+        )
+
+    prev_ctx_len = 0
+    for r in range(num_rounds):
+        tokens = torch.randint(0, config.vocab_size, (1, block_len), dtype=torch.long, generator=gen)
+        S = ctx.shape[1]
+        with torch.no_grad():
+            if use_cache:
+                new_ctx_len = S - prev_ctx_len
+                cache_position = torch.arange(prev_ctx_len, S, dtype=torch.long)
+                logits = model(
+                    tokens, ctx,
+                    cache=cache, cache_position=cache_position, new_ctx_len=new_ctx_len,
+                )
+                prev_ctx_len = S
+            else:
+                logits = model(tokens, ctx)
+        per_round_logits.append(logits)
+
+        accepted = (r % block_len) + 1
+        new_ctx = torch.randn(1, accepted, hidden_size, generator=gen)
+        ctx = torch.cat([ctx, new_ctx], dim=1)
+
+    return per_round_logits
+
+
+def main():
+    model, config = make_model()
+
+    reference_logits = simulate_rounds(model, config, use_cache=False)
+    print(f"Captured uncached reference for {len(reference_logits)} rounds.")
+
+    cached_logits = simulate_rounds(model, config, use_cache=True)
+    print(f"Captured cached path for {len(cached_logits)} rounds.")
+
+    max_diff = max(
+        (a - b).abs().max().item()
+        for a, b in zip(reference_logits, cached_logits)
+    )
+    print(f"\ncached-vs-uncached max abs diff across rounds: {max_diff:.3e}")
+    for i, (a, b) in enumerate(zip(reference_logits, cached_logits)):
+        d = (a - b).abs().max().item()
+        print(f"  round {i}: diff {d:.3e}")
+
+    if max_diff < 1e-4:
+        print("\nPASS: cached path matches uncached -- wiring is numerically correct.")
+    else:
+        print("\nFAIL: cached path diverges from uncached -- wiring bug.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scratch_draft_cache_probe.py b/scratch_draft_cache_probe.py
new file mode 100644
index 00000000000..58160ea5a07
--- /dev/null
+++ b/scratch_draft_cache_probe.py
@@ -0,0 +1,84 @@
+"""Standalone probe: does DFlashDraftKVCache correctly self-heal across
+rounds, the way the design walkthrough claimed?
+
+Not part of the PR -- a throwaway diagnostic, same style as
+scratch_cache_probe.py earlier. Run directly:
+    python scratch_draft_cache_probe.py
+
+Mirrors the concrete walkthrough numbers from the design discussion:
+    Round 1: ctx_len starts at 20 (a 20-token prompt). Block length 8.
+             Block written at [20, 28).
+    Round 2: 5 of those 8 accepted -> ctx_len becomes 25. New block (8 more)
+             written at [25, 33) -- overwriting round 1's now-stale [25, 28)
+             tail and extending past it.
+
+If this passes, positions 20-24 should still hold round 1's block content
+(the part that got confirmed), and everything from 25 onward should be
+entirely round 2's content, with zero leftover from round 1's rejected
+speculative tail.
+"""
+
+import torch
+
+from executorch.backends.mlx.examples.llm.dflash_draft_cache import DFlashDraftKVCache
+
+
+def fake_kv(num_heads, head_dim, seq_len, fill_value):
+    shape = (1, num_heads, seq_len, head_dim)
+    return (
+        torch.full(shape, fill_value, dtype=torch.float32),
+        torch.full(shape, fill_value, dtype=torch.float32),
+    )
+
+
+def main():
+    num_layers, num_heads, head_dim, max_seq_len = 1, 2, 4, 64
+    cache = DFlashDraftKVCache(num_layers, num_heads, head_dim, max_seq_len)
+
+    # Simulate a prompt already having filled positions [0, 20) with some
+    # earlier value (3.0), just so there's realistic prior content.
+    k0, v0 = fake_kv(num_heads, head_dim, 20, fill_value=3.0)
+    cache.write_context(0, k0, v0, start_pos=0)
+    cache.advance_context(torch.tensor(20, dtype=torch.long))
+
+    # --- Round 1: block of 8 written at [20, 28), filled with 1.0 ---
+    k1, v1 = fake_kv(num_heads, head_dim, 8, fill_value=1.0)
+    start_pos_1 = cache.ctx_len.item()
+    assert start_pos_1 == 20, f"expected ctx_len=20 before round 1, got {start_pos_1}"
+    cache.write_block(0, k1, v1, start_pos=start_pos_1)
+    print(f"After round 1 block write, k_cache[18:30]: "
+          f"{cache.layers[0].k_cache[0, 0, 18:30, 0].tolist()}")
+
+    # --- Round 2: only 5 of round 1's 8 accepted -> confirmed context grows
+    # by 5 (not 8). New block (8 more) written starting at the new ctx_len. ---
+    k_confirmed, v_confirmed = fake_kv(num_heads, head_dim, 5, fill_value=1.0)
+    start_pos_confirm = cache.ctx_len.item()
+    assert start_pos_confirm == 20
+    cache.write_context(0, k_confirmed, v_confirmed, start_pos=start_pos_confirm)
+    cache.advance_context(torch.tensor(5, dtype=torch.long))
+
+    start_pos_2 = cache.ctx_len.item()
+    assert start_pos_2 == 25, f"expected ctx_len=25 before round 2 block, got {start_pos_2}"
+    k2, v2 = fake_kv(num_heads, head_dim, 8, fill_value=2.0)
+    cache.write_block(0, k2, v2, start_pos=start_pos_2)
+
+    result = cache.layers[0].k_cache[0, 0, 18:33, 0].tolist()
+    print(f"After round 2 block write, k_cache[18:33]: {result}")
+
+    expected = [3.0, 3.0] + [1.0] * 5 + [2.0] * 8
+    if result == expected:
+        print("PASS: self-healing overwrite worked exactly as designed.")
+    else:
+        print(f"FAIL: expected {expected}, got {result}")
+
+    # Sanity check the mask too: with ctx_len=25 and this round's block_len=8,
+    # valid range should be exactly [0, 33).
+    mask = cache.valid_mask(block_len=8)
+    valid_count = mask.sum().item()
+    print(f"valid_mask() reports {valid_count} valid positions (expected 33).")
+    assert valid_count == 33, f"expected 33 valid positions, got {valid_count}"
+    print("PASS: valid_mask boundary is correct.")
+
+
+if __name__ == "__main__":
+    main()