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/examples/llm/dflash_draft_model.py b/backends/mlx/examples/llm/dflash_draft_model.py new file mode 100644 index 00000000000..3562d840b72 --- /dev/null +++ b/backends/mlx/examples/llm/dflash_draft_model.py @@ -0,0 +1,275 @@ +# 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), 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 +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 + # 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: + # 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) + ) + 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 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): + # 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, :, :] + 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): + """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.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): + """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( + 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 = 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 + ) + 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, :] + 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): + # 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)) + + +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) + # 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) + # 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"), + ) diff --git a/backends/mlx/examples/llm/export_llm_hf.py b/backends/mlx/examples/llm/export_llm_hf.py index fe6b8094f6b..91b081696c8 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.examples.models.qwen3.mlx_source_transformations 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/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 new file mode 100644 index 00000000000..a3c18866764 --- /dev/null +++ b/examples/models/qwen3/export_dflash_draft.py @@ -0,0 +1,117 @@ +# 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(): + 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) + 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) + position_ids = torch.arange(ctx_len + block_size).unsqueeze(0) + + 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.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}, block_size fixed at {block_size}." + ) + + +if __name__ == "__main__": + 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..16a293c34e7 --- /dev/null +++ b/examples/models/qwen3/mlx_source_transformations.py @@ -0,0 +1,69 @@ +# 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 -- +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 +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 +): + + 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/examples/models/qwen3/run_baseline.py b/examples/models/qwen3/run_baseline.py new file mode 100644 index 00000000000..594dada473a --- /dev/null +++ b/examples/models/qwen3/run_baseline.py @@ -0,0 +1,77 @@ +# 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..d4c683ea86d --- /dev/null +++ b/examples/models/qwen3/run_dflash.py @@ -0,0 +1,217 @@ +# 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, , , ...] + 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).", + ) + 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 + + # 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 + # 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. 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, + ) + 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_exec_time = time.time() - _t0 + draft_ids = draft_logits[0].argmax(-1).tolist() # block_size - 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() # block_size 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]] + accepted_total += accepted + emitted_total += len(new_tokens) + + # 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] + + 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}") + 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_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)}")