diff --git a/.gitignore b/.gitignore index ee206e23d94..0ef50c156a2 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,9 @@ zephyr_dev_root.backup.*/ # Agents .claude/*.local.* extension/pybindings/mlx.metallib + +# Scratch/WIP work not ready for review -- never committed +/wip/ +gemma-4-31B-it-DFlash/ +gemma-4-31B-it-HQQ-INT4/ +gemma4_31b_dflash_exports_mlx/ 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/CMakeLists.txt b/backends/mlx/CMakeLists.txt index c6eb2ac3156..409c390e458 100644 --- a/backends/mlx/CMakeLists.txt +++ b/backends/mlx/CMakeLists.txt @@ -186,7 +186,7 @@ message( # each patch file under patches/ for its rationale. set(_mlx_patches ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_nax_has_include.patch - ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_qmm_splitk_bk_align.patch + # ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_qmm_splitk_bk_align.patch # Disabled: targets qmm_splitk, introduced in MLX v0.32.0 (this repo is pinned to an earlier commit); not applicable, see patch file header ) ExternalProject_Add( mlx_external 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..6cd3debbdf2 --- /dev/null +++ b/backends/mlx/examples/llm/dflash_draft_model.py @@ -0,0 +1,277 @@ +# 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 + # Upcast to float32 for the norm (numerical stability), matching the + # previous manual pow/mean/rsqrt implementation's order of operations + # exactly: normalize in fp32, downcast, THEN multiply by weight. + x = torch.nn.functional.rms_norm(x.float(), (x.shape[-1],), eps=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/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/patches/mlx_nax_has_include.patch b/backends/mlx/patches/mlx_nax_has_include.patch index 2c917353146..0a896e6f4ba 100644 --- a/backends/mlx/patches/mlx_nax_has_include.patch +++ b/backends/mlx/patches/mlx_nax_has_include.patch @@ -1,43 +1,49 @@ Guard the NAX MetalPerformancePrimitives includes behind __has_include. MLX's NAX kernels (GEMM and attention) unconditionally include -, a framework that -only ships in the macOS 26 / Xcode 26 SDK. On older SDKs the JIT preamble -generator (make_compiled_preamble.sh) runs `metal -E` over these headers, which -fatals on the missing include and fails the build. MLX already gates NAX on the -non-JIT metallib path (kernels/CMakeLists.txt), but the JIT path -(MLX_METAL_JIT=ON, which ExecuTorch uses) is ungated. +, a framework +that only ships in the macOS 26 / Xcode 26 SDK. On older SDKs the JIT +preamble generator (make_compiled_preamble.sh) runs `metal -E` over +these headers, which fatals on the missing include and fails the build. +MLX already gates NAX on the non-JIT metallib path (kernels/CMakeLists.txt), +but the JIT path (MLX_METAL_JIT=ON, which ExecuTorch uses) is ungated. -This wraps the includes in `#if __has_include(...)` so preprocessing succeeds on -pre-26 SDKs (NAX kernels are never JIT-compiled at runtime there anyway, since -device.cpp:is_nax_available() returns false), while macOS 26 still gets NAX. +This wraps the includes in `#if __has_include(...)` so preprocessing +succeeds on pre-26 SDKs (NAX kernels are never JIT-compiled at runtime +there anyway, since device.cpp:is_nax_available() returns false), while +macOS 26 still gets NAX. Upstream candidate; carried locally until an MLX release gates the JIT path. +Regenerated against the currently-pinned MLX submodule commit: gemm/nax.h +picked up an extra `#include ".../gemm/transforms.h"` line upstream since +this patch was first written, shifting its context. attn/nax.h's context +was unchanged. + diff --git a/mlx/backend/metal/kernels/steel/gemm/nax.h b/mlx/backend/metal/kernels/steel/gemm/nax.h +index 5839176..05f3320 100644 --- a/mlx/backend/metal/kernels/steel/gemm/nax.h +++ b/mlx/backend/metal/kernels/steel/gemm/nax.h -@@ -9,7 +9,9 @@ - #include "mlx/backend/metal/kernels/steel/defines.h" +@@ -10,7 +10,9 @@ + #include "mlx/backend/metal/kernels/steel/gemm/transforms.h" #include "mlx/backend/metal/kernels/steel/utils/integral_constant.h" --#include +#if __has_include() -+#include + #include +#endif using namespace metal; diff --git a/mlx/backend/metal/kernels/steel/attn/nax.h b/mlx/backend/metal/kernels/steel/attn/nax.h +index c8f3ea5..3e9542d 100644 --- a/mlx/backend/metal/kernels/steel/attn/nax.h +++ b/mlx/backend/metal/kernels/steel/attn/nax.h @@ -9,7 +9,9 @@ #include "mlx/backend/metal/kernels/steel/defines.h" #include "mlx/backend/metal/kernels/steel/utils/integral_constant.h" --#include +#if __has_include() -+#include + #include +#endif using namespace metal; diff --git a/backends/mlx/third-party/mlx b/backends/mlx/third-party/mlx index 7a1d4f5c12a..ce45c52505c 160000 --- a/backends/mlx/third-party/mlx +++ b/backends/mlx/third-party/mlx @@ -1 +1 @@ -Subproject commit 7a1d4f5c12ac82f4b4d0a6e71538d89ca0605247 +Subproject commit ce45c52505c8158ea48d2a54e8caae05efd86bfe diff --git a/examples/models/gemma4_31b/CMakeLists.txt b/examples/models/gemma4_31b/CMakeLists.txt index c740de13c36..cca4e87e81a 100644 --- a/examples/models/gemma4_31b/CMakeLists.txt +++ b/examples/models/gemma4_31b/CMakeLists.txt @@ -74,6 +74,14 @@ target_include_directories( ) target_link_libraries(gemma4_31b_worker PUBLIC ${link_libraries}) +# Standalone C++ DFlash driver -- benchmarks speculative decoding via +# Module API directly, bypassing the Python pybind runtime entirely. +add_executable(dflash_cpp_driver dflash_cpp_driver.cpp) +target_include_directories( + dflash_cpp_driver PUBLIC ${_common_include_directories} ${_json_include} +) +target_link_libraries(dflash_cpp_driver PUBLIC ${link_libraries}) + if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug") target_link_options_gc_sections(gemma4_31b_runner) if(NOT APPLE AND NOT MSVC) @@ -88,6 +96,7 @@ endif() if(TARGET mlxdelegate) executorch_target_copy_mlx_metallib(gemma4_31b_runner) executorch_target_copy_mlx_metallib(gemma4_31b_worker) + executorch_target_copy_mlx_metallib(dflash_cpp_driver) endif() if(EXECUTORCH_BUILD_CUDA) diff --git a/examples/models/gemma4_31b/DFLASH_EXPERIMENTS.md b/examples/models/gemma4_31b/DFLASH_EXPERIMENTS.md new file mode 100644 index 00000000000..e3db3cad2f3 --- /dev/null +++ b/examples/models/gemma4_31b/DFLASH_EXPERIMENTS.md @@ -0,0 +1,117 @@ +Written By: Chetan Thotti (cthotti) +Date: 07/24/2026 + +This is a record of DFlash speculative decoding benchmarking for +Gemma4-31B on a single M4 Pro machine (10-core CPU / 20-core GPU / 64GB +LPDDR5X). The short version: **DFlash gives a real, solid speedup on +math and code prompts (~1.5-1.7x), but is slower than plain baseline +decoding on open-ended chat prompts (~0.77x).** The task-dependence is +the headline finding here, not the raw tok/s numbers -- see the Qwen3 +write-up (`../qwen3/DFLASH_EXPERIMENTS.md`) for the hardware-generation +story, which we didn't re-run here (only one machine tested). + +## The setup + +Target model was the prequantized `SocialLocalMobile/gemma-4-31B-it-HQQ-INT4` +checkpoint (INT4/INT8 mixed, ~21GB, "sensitive" recipe), draft was +`z-lab/gemma-4-31B-it-DFlash` (5 layers, 6 tapped target layers +`[1,12,23,35,46,57]`, block_size 16). Both baseline and DFlash runs reuse +the same target `.pte` (the one with the two-output `(logits, hidden)` +signature) rather than a separately-exported plain baseline model -- same +methodology as the Qwen3 write-up, so the comparison isolates the +algorithm rather than any difference between two export paths. + +Getting a correct target export here took two real fixes, not just a +straightforward port from Qwen3: + +- The base `Gemma4_31B.forward` (and the MLX-optimized version installed + by `mlx_source_transformations.py`) is last-token-only by design -- + correct for normal decode, wrong for DFlash verification, which needs + logits at every position in the drafted block. Using it unmodified + silently collapsed the exported model back to `(B, 1, V)` with only + one output, dropping the hidden-state output entirely, no error at any + stage. +- `mlx_source_transformations.py`'s `_replace_model_forward` does a full + `types.MethodType` replacement of the top-level forward at export time + -- it doesn't wrap or call whatever forward the model already had, it + discards it. So even after fixing `Gemma4_31BWithHidden.forward` + directly, calling the stock transform function still silently reverted + the model to last-token-only, single-output behavior. Needed a + Gemma4-specific `dflash_mlx_source_transformations.py` that reuses the + per-layer MLX op rewrites (rope/custom_sdpa/kv-cache, which are + output-shape-agnostic) but installs a DFlash-aware top-level forward + instead of the stock one. + +Both were verified by checking the exported `.pte`'s actual `MethodMeta` +(`num_outputs`, tensor shapes) after export, not just that export +completed without error -- the failure mode here doesn't throw. + +## Results + +Three prompts (math, code, chat), three trials each, 300 max new tokens, +greedy decoding: + +| Category | Baseline tok/s | DFlash tok/s | tau | Speedup | +|----------|---------------|--------------|------|---------| +| Math | 10.41 | 15.23 | 6.52 | 1.46x | +| Code | 10.22 | 17.58 | 7.70 | 1.72x | +| Chat | 10.48 | 8.02 | 3.48 | 0.77x | + +(Prompts: math was a word problem requiring algebraic setup and solving; +code was "write binary search + explain time complexity"; chat was +"explain how photosynthesis works in detail". Within-category trials +were tightly consistent -- greedy decoding is deterministic, so this is +mostly a sanity check that nothing was flaky, not real independent +samples.) + +## Why math/code work and chat doesn't + +Target verification cost (`target_exec`) is close to flat regardless of +context length -- roughly 374ms per round whether verifying a 20-token or +90-token window. That's the expected signature of a memory-bandwidth- +bound model: verifying a block of drafted tokens in one forward pass +costs close to the same as verifying one token, because the dominant +cost is streaming the ~21GB of weights through memory once, not the +compute itself. That flatness is what makes DFlash's speedup possible at +all on this hardware -- the breakeven tau (where DFlash stops being a +net loss) works out to roughly 3.8-3.9 given baseline's ~10.3 tok/s and +target_exec's ~374ms. + +Math and code prompts land comfortably above that line (tau 6.5-7.7). +Chat lands just below it (tau 3.48), which is why it's the one category +that's slower than baseline rather than faster. This mirrors the DFlash +paper's own reported numbers almost exactly: their Table 1 shows +MT-Bench (chat/conversational) as the consistently weakest category +across every model they tested, well below math/code, even on much +smaller target models and H200 GPUs. So this isn't a Gemma4-specific +weakness or a bug in our export -- it's the same task-dependent pattern +the paper documents, just reproduced on a bigger target model and +different hardware. + +## What we didn't get to + +- **Only one machine tested** (M4 Pro). The Qwen3 write-up found DFlash's + speedup depends heavily on GPU architecture generation (Apple9/M3+ + vs Apple8/M1-M2), not on raw core count or bandwidth. We'd expect the + same story to hold for Gemma4 -- an M2-class chip likely wouldn't show + a speedup even on math/code -- but that's inference from the Qwen3 + result, not something confirmed on Gemma4 directly. +- **An incremental-caching optimization for the draft model's context + K/V was attempted and reverted.** The draft model recomputes + projections/norm/RoPE over the *entire* accumulated target-hidden + history every round rather than just the new chunk since the last + round -- real, measurable waste that gets worse as generations get + longer. We built and numerically validated (float32-precision-exact) + an incremental-caching version, but hit a `torch.export` lowering + failure (`GuardOnDataDependentSymNode` on a dynamic slice of an + unbacked SymInt) that would have needed a deeper rewrite to + `mlx::custom_sdpa` to fix properly. Separately, and more importantly: + the internal DFlash-on-ExecuTorch design doc states the draft cache is + *intentionally* reset every round ("the draft model re-processes + context hidden states each round, so rollback is not needed") -- so + the behavior we were trying to optimize away may be the intended + reference design, not an oversight. Reverted rather than pushing + further against that. +- **MLX v0.32.0's small-M GEMM improvements** (flagged in the Qwen3 PR + review) were never tested against Gemma4 either -- this repo is still + pinned to v0.31.1. diff --git a/examples/models/gemma4_31b/dflash_cpp_driver.cpp b/examples/models/gemma4_31b/dflash_cpp_driver.cpp new file mode 100644 index 00000000000..3b909ea92b5 --- /dev/null +++ b/examples/models/gemma4_31b/dflash_cpp_driver.cpp @@ -0,0 +1,325 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * Standalone C++ DFlash speculative decoding driver for Gemma4-31B. + * Ports run_dflash.py's draft/verify/accept loop directly using + * ExecuTorch's Module API (raw tensor I/O), since the higher-level + * LLMEngine/LLMSession framework used by main.cpp is built for + * single-model token-at-a-time generation, not this two-model + * draft-then-verify pattern with multi-output tensors. + */ + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using executorch::extension::Module; +using executorch::extension::make_tensor_ptr; +using executorch::extension::TensorPtr; +using executorch::runtime::EValue; +using executorch::runtime::Error; + +DEFINE_string(target_pte, "gemma4_31b_dflash_exports_mlx/model.pte", "Target .pte path."); +DEFINE_string(draft_pte, "gemma4_31b_dflash_draft.pte", "Draft .pte path."); +DEFINE_string(tokenizer_path, "gemma-4-31B-it-HQQ-INT4/tokenizer.json", "Tokenizer path."); +DEFINE_string(prompt, "The capital of France is", "Prompt text."); +DEFINE_int32(max_new_tokens, 64, "Maximum tokens to generate."); +DEFINE_int32(block_size, 16, "DFlash draft block size."); +DEFINE_int32(mask_id, 4, "DFlash mask token id (from draft checkpoint config)."); +DEFINE_bool(raw_prompt, false, "Skip chat-template wrapping."); +DEFINE_bool(verbose, false, "Print per-round timing/acceptance debug output."); + +namespace { + +constexpr int64_t kBosId = 2; +const std::set kEosIds = {1, 50, 106}; + +std::string format_prompt(const std::string& prompt) { + return "<|turn>user\n" + prompt + + "\n<|turn>model\n<|channel>thought\n"; +} + +double now_ms() { + return std::chrono::duration( + std::chrono::steady_clock::now().time_since_epoch()) + .count(); +} + +std::vector argmax_last_dim(const executorch::aten::Tensor& t) { + auto sizes = t.sizes(); + int64_t T = sizes[1]; + int64_t V = sizes[2]; + const float* data = t.const_data_ptr(); + std::vector out(T); + for (int64_t i = 0; i < T; ++i) { + const float* row = data + i * V; + out[i] = std::max_element(row, row + V) - row; + } + return out; +} + +int64_t first_mismatch( + const std::vector& draft_ids, + const std::vector& target_ids) { + for (size_t i = 0; i < draft_ids.size(); ++i) { + if (draft_ids[i] != target_ids[i]) { + return static_cast(i); + } + } + return static_cast(draft_ids.size()); +} + +// The target's hidden-state output is BFloat16 (confirmed via MethodMeta: +// dtype=BFloat16), NOT Float32 -- treating its raw bytes as float32 via +// const_data_ptr() silently reinterprets 2-byte bf16 values as +// 4-byte floats, producing garbage that made the draft's context +// essentially random (observed: near-zero acceptance, ctx_len advancing +// by ~1 per round instead of the expected 5-15). bf16 is simply the top +// 16 bits of a float32 (same exponent width, truncated mantissa), so +// converting is a left-shift into the upper half of a 32-bit word. +std::vector bf16_tensor_to_fp32(const executorch::aten::Tensor& t) { + int64_t numel = t.numel(); + const uint16_t* src = reinterpret_cast(t.const_data_ptr()); + std::vector out(numel); + for (int64_t i = 0; i < numel; ++i) { + uint32_t bits = static_cast(src[i]) << 16; + std::memcpy(&out[i], &bits, sizeof(float)); + } + return out; +} + +} // namespace + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + + auto tokenizer = std::make_unique(); + if (tokenizer->load(FLAGS_tokenizer_path) != tokenizers::Error::Ok) { + ET_LOG(Error, "Failed to load tokenizer: %s", FLAGS_tokenizer_path.c_str()); + return 1; + } + + Module target_module(FLAGS_target_pte); + Module draft_module(FLAGS_draft_pte); + if (target_module.load() != Error::Ok) { + ET_LOG(Error, "Failed to load target .pte: %s", FLAGS_target_pte.c_str()); + return 1; + } + if (draft_module.load() != Error::Ok) { + ET_LOG(Error, "Failed to load draft .pte: %s", FLAGS_draft_pte.c_str()); + return 1; + } + + std::string prompt_text = + FLAGS_raw_prompt ? FLAGS_prompt : format_prompt(FLAGS_prompt); + auto encoded = tokenizer->encode(prompt_text, /*bos=*/0, /*eos=*/0); + if (!encoded.ok()) { + ET_LOG(Error, "Failed to encode prompt"); + return 1; + } + std::vector prompt_ids; + prompt_ids.push_back(kBosId); + for (auto id : encoded.get()) { + prompt_ids.push_back(static_cast(id)); + } + int64_t prompt_len = static_cast(prompt_ids.size()); + printf("Prompt tokens: %lld\n", (long long)prompt_len); + + std::vector input_pos_vec(prompt_len); + for (int64_t i = 0; i < prompt_len; ++i) input_pos_vec[i] = i; + + auto prompt_ids_tensor = + make_tensor_ptr({1, (int)prompt_len}, prompt_ids.data(), executorch::aten::ScalarType::Long); + auto input_pos_tensor = make_tensor_ptr( + {(int)prompt_len}, input_pos_vec.data(), executorch::aten::ScalarType::Long); + + auto prefill_result = + target_module.forward({EValue(prompt_ids_tensor), EValue(input_pos_tensor)}); + if (!prefill_result.ok()) { + ET_LOG(Error, "Prefill forward failed"); + return 1; + } + auto prefill_outputs = prefill_result.get(); + auto logits = prefill_outputs[0].toTensor(); + auto hidden = prefill_outputs[1].toTensor(); + + int64_t pos = prompt_len; + auto logits_argmax = argmax_last_dim(logits); + int64_t last_token = logits_argmax.back(); + + std::vector generated = {last_token}; + int64_t rounds = 0; + int64_t accepted_total = 0; + int64_t emitted_total = 0; + + + + int64_t hidden_concat_dim = hidden.sizes()[2]; + std::vector hidden_history = bf16_tensor_to_fp32(hidden); + int64_t hidden_len = prompt_len; + + // Warm-up: MLX/Metal lazily JIT-compiles kernels on a delegate's FIRST + // real execution, not at Module::load() time. Without this, that one-time + // compile cost lands inside round 1's timed draft_exec (observed: ~296ms + // vs ~40ms steady-state) and skews the reported tokens/s. Run one + // throwaway forward with the actual round-1 inputs, discard the result, + // so the timed loop below only measures steady-state performance -- + // matching what Python's Runtime.load_program().load_method() chain + // appears to already do more eagerly at load time. + { + std::vector warm_input_vec(FLAGS_block_size); + warm_input_vec[0] = last_token; + for (int64_t i = 1; i < FLAGS_block_size; ++i) warm_input_vec[i] = FLAGS_mask_id; + auto warm_input_tensor = make_tensor_ptr( + {1, (int)FLAGS_block_size}, warm_input_vec.data(), executorch::aten::ScalarType::Long); + auto warm_hidden_tensor = make_tensor_ptr( + {1, (int)hidden_len, (int)hidden_concat_dim}, + hidden_history.data(), + executorch::aten::ScalarType::Float); + std::vector warm_pos_vec(hidden_len + FLAGS_block_size); + for (int64_t i = 0; i < hidden_len + FLAGS_block_size; ++i) warm_pos_vec[i] = i; + auto warm_pos_tensor = make_tensor_ptr( + {1, (int)(hidden_len + FLAGS_block_size)}, + warm_pos_vec.data(), + executorch::aten::ScalarType::Long); + auto warm_result = draft_module.forward( + {EValue(warm_input_tensor), EValue(warm_hidden_tensor), EValue(warm_pos_tensor)}); + if (!warm_result.ok()) { + ET_LOG(Error, "Draft warm-up forward failed"); + return 1; + } + } + + // Timer starts AFTER prefill AND warm-up, matching run_dflash.py's t0 + // placement -- tokens/s measures only the speculative decoding round + // loop's steady-state performance. + double t0 = now_ms(); + + while (static_cast(generated.size()) < FLAGS_max_new_tokens) { + rounds++; + int64_t bs = FLAGS_block_size; + + std::vector draft_input_vec(bs); + draft_input_vec[0] = last_token; + for (int64_t i = 1; i < bs; ++i) draft_input_vec[i] = FLAGS_mask_id; + auto draft_input_tensor = + make_tensor_ptr({1, (int)bs}, draft_input_vec.data(), executorch::aten::ScalarType::Long); + + auto hidden_tensor = make_tensor_ptr( + {1, (int)hidden_len, (int)hidden_concat_dim}, + hidden_history.data(), + executorch::aten::ScalarType::Float); + + std::vector draft_pos_vec(hidden_len + bs); + for (int64_t i = 0; i < hidden_len + bs; ++i) draft_pos_vec[i] = i; + auto draft_pos_tensor = make_tensor_ptr( + {1, (int)(hidden_len + bs)}, draft_pos_vec.data(), executorch::aten::ScalarType::Long); + + double dt0 = now_ms(); + auto draft_result = draft_module.forward( + {EValue(draft_input_tensor), EValue(hidden_tensor), EValue(draft_pos_tensor)}); + double draft_exec_ms = now_ms() - dt0; + if (!draft_result.ok()) { + ET_LOG(Error, "Draft forward failed at round %lld", (long long)rounds); + return 1; + } + auto draft_logits = draft_result.get()[0].toTensor(); + auto draft_ids = argmax_last_dim(draft_logits); + + std::vector verify_input_vec; + verify_input_vec.push_back(last_token); + verify_input_vec.insert(verify_input_vec.end(), draft_ids.begin(), draft_ids.end()); + int64_t verify_len = static_cast(verify_input_vec.size()); + auto verify_input_tensor = make_tensor_ptr( + {1, (int)verify_len}, verify_input_vec.data(), executorch::aten::ScalarType::Long); + + std::vector verify_pos_vec(verify_len); + for (int64_t i = 0; i < verify_len; ++i) verify_pos_vec[i] = pos + i; + auto verify_pos_tensor = make_tensor_ptr( + {(int)verify_len}, verify_pos_vec.data(), executorch::aten::ScalarType::Long); + + double vt0 = now_ms(); + auto verify_result = target_module.forward( + {EValue(verify_input_tensor), EValue(verify_pos_tensor)}); + double target_exec_ms = now_ms() - vt0; + if (!verify_result.ok()) { + ET_LOG(Error, "Verify forward failed at round %lld", (long long)rounds); + return 1; + } + auto verify_outputs = verify_result.get(); + auto target_logits = verify_outputs[0].toTensor(); + auto new_hidden = verify_outputs[1].toTensor(); + auto target_ids = argmax_last_dim(target_logits); + + int64_t accepted = first_mismatch(draft_ids, target_ids); + if (FLAGS_verbose && rounds <= 10) { + printf( + " timing: draft_exec=%.1fms target_exec=%.1fms ctx_len=%lld\n", + draft_exec_ms, + target_exec_ms, + (long long)hidden_len); + } + + std::vector new_tokens(draft_ids.begin(), draft_ids.begin() + accepted); + new_tokens.push_back(target_ids[accepted]); + + bool hit_eos = false; + for (size_t i = 0; i < new_tokens.size(); ++i) { + if (kEosIds.count(new_tokens[i])) { + new_tokens.resize(i + 1); + accepted = std::min(accepted, static_cast(new_tokens.size()) - 1); + hit_eos = true; + break; + } + } + + accepted_total += accepted; + emitted_total += static_cast(new_tokens.size()); + generated.insert(generated.end(), new_tokens.begin(), new_tokens.end()); + + pos += static_cast(new_tokens.size()); + last_token = new_tokens.back(); + int64_t append_len = static_cast(new_tokens.size()); + std::vector new_hidden_fp32 = bf16_tensor_to_fp32(new_hidden); + hidden_history.insert( + hidden_history.end(), + new_hidden_fp32.begin(), + new_hidden_fp32.begin() + append_len * hidden_concat_dim); + hidden_len += append_len; + + if (hit_eos) break; + } + + double total_ms = now_ms() - t0; + int64_t n = static_cast(generated.size()); + + printf("\nPrompt: %s\n", FLAGS_prompt.c_str()); + printf("Generated (%lld tokens)\n", (long long)n); + printf("\n--stats--\n"); + printf("rounds: %lld\n", (long long)rounds); + printf( + "avg accepted/round (draft-only): %.2f\n", + static_cast(accepted_total) / rounds); + printf( + "avg emitted/round (tau, incl. bonus): %.2f\n", + static_cast(emitted_total) / rounds); + printf( + "time: %.2fs tokens/s: %.2f\n", + total_ms / 1000.0, + n / (total_ms / 1000.0)); + + return 0; +} diff --git a/examples/models/gemma4_31b/dflash_hidden_export.py b/examples/models/gemma4_31b/dflash_hidden_export.py new file mode 100644 index 00000000000..dfb5892d6ff --- /dev/null +++ b/examples/models/gemma4_31b/dflash_hidden_export.py @@ -0,0 +1,84 @@ +# 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. + +"""Gemma4-31B DFlash hidden-state export wrapper. + +Gemma4-31B's forward() is hand-written (not the generic export_llm_hf.py +path), so backends/mlx/examples/llm/dflash_hidden_export.py's +TorchExportableModuleWithStaticCacheAndHidden doesn't apply here. This +patches Gemma4_31B.forward directly to also return hidden states from +the configured target layers. + +target_layer_ids are 0-indexed into self.layers (captured[i] = output of +self.layers[i]), matching the Qwen3 wrapper's convention. + +UNVERIFIED: Gemma4DecoderLayer applies layer_scalar before returning, so +the captured hidden already includes it. Not confirmed against HF's +modeling_gemma4.py (unavailable in current transformers) or z-lab's +training code -- if wrong, draft conditioning will be subtly off and +tau will look worse with no explicit error. Re-check before trusting +benchmarks on this path. +""" + +from typing import List, Sequence, Tuple + +import torch + +from executorch.examples.models.gemma4_31b.model import Gemma4_31B, Gemma4_31BConfig + + +class Gemma4_31BWithHidden(Gemma4_31B): + """Gemma4_31B variant that also returns concatenated target hidden states.""" + + def __init__(self, config: Gemma4_31BConfig, layer_ids: Sequence[int] = ()): + super().__init__(config) + if not layer_ids: + raise ValueError("layer_ids must be non-empty") + self.dflash_layer_ids: List[int] = list(layer_ids) + + def forward( + self, + tokens: torch.LongTensor, + input_pos: torch.LongTensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Full-sequence logits + hidden states for DFlash block verification. + + Unlike the base Gemma4_31B.forward (last-logits-only + on-device + sampling -- correct for single-token decode), DFlash verifies an + entire drafted block (T positions) in one forward pass, comparing + the target's argmax at each position against the draft's proposed + token there. Last-token-only logits would silently drop every + position but the final one, breaking block verification. Sampling + is left to the caller (run_dflash.py does greedy argmax / + first-mismatch there), so no temperature argument or sample() call. + + Returns: + logits: (B, T, vocab_size) float, softcapped, all positions. + hidden: (B, T, len(dflash_layer_ids) * hidden_size). + """ + x = self.embed_tokens(tokens) * self.embed_normalizer + sliding_mask, full_mask = self._build_masks(input_pos) + + layer_id_set = set(self.dflash_layer_ids) + captured = {} + for i, layer in enumerate(self.layers): + x = layer(x, input_pos, sliding_mask, full_mask) + if i in layer_id_set: + captured[i] = x + + missing = layer_id_set - captured.keys() + if missing: + raise ValueError( + f"dflash_layer_ids {sorted(missing)} not reached -- " + f"model only has {len(self.layers)} layers" + ) + hidden = torch.cat([captured[i] for i in self.dflash_layer_ids], dim=-1) + + x = self.norm(x) + logits = self.lm_head(x).float() # (B, T, vocab) -- NOT x[:, -1, :] + cap = self.logit_softcap.float() + logits = torch.tanh(logits / cap) * cap + return logits, hidden diff --git a/examples/models/gemma4_31b/dflash_mlx_source_transformations.py b/examples/models/gemma4_31b/dflash_mlx_source_transformations.py new file mode 100644 index 00000000000..b272db63951 --- /dev/null +++ b/examples/models/gemma4_31b/dflash_mlx_source_transformations.py @@ -0,0 +1,147 @@ +# 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. + +"""DFlash-aware MLX source transformations for Gemma4-31B. + +mlx_source_transformations.py's mlx_source_transformations() ends by calling +_replace_model_forward(model), which unconditionally overwrites the top-level +forward with a (tokens, input_pos) -> (B, 1, V) last-token-only, single-output +variant -- it does a full `types.MethodType` replacement, not a wrapper, so +whatever forward the model had before (including Gemma4_31BWithHidden's +full-sequence, two-output forward) is discarded entirely and never called at +export time. That silently turned our DFlash target export back into a +last-token-only, hidden-state-free model with no error at any stage. + +This module reuses everything else from mlx_source_transformations.py +unchanged (KV cache attachment, per-layer attention/layer forward rewrites +via mlx.rope / mlx.custom_sdpa -- both output-shape-agnostic, operate within +a layer) and only replaces the final _replace_model_forward(model) call with +a DFlash-aware top-level forward that preserves full-sequence logits and +hidden-state capture at model.dflash_layer_ids. +""" + +import types + +import torch +import torch.nn as nn + +from executorch.examples.models.gemma4_31b.mlx_source_transformations import ( + MLXKVCache, + MLXRingKVCache, + MLXTurboQuantKVCache, + _replace_attention_forward, + _replace_layer_forward, +) + + +def _replace_dflash_model_forward(model: nn.Module) -> None: + """Replace the top-level forward with a DFlash-aware, MLX-optimized variant. + + Signature: (tokens, input_pos) -> (logits, hidden) where logits is + (B, T, V) over every position (not last-token-only) and hidden is + (B, T, len(dflash_layer_ids) * hidden_size) -- matching + Gemma4_31BWithHidden.forward's contract, but using the per-layer + MLX-optimized attention/layer forwards installed by + dflash_mlx_source_transformations instead of the original PyTorch ops. + """ + + def _mlx_dflash_model_forward( + self, tokens: torch.Tensor, input_pos: torch.Tensor + ): + x = self.embed_tokens(tokens) * self.embed_normalizer + + layer_id_set = set(self.dflash_layer_ids) + captured = {} + for i, layer in enumerate(self.layers): + x = layer(x, input_pos) + if i in layer_id_set: + captured[i] = x + + missing = layer_id_set - captured.keys() + if missing: + raise ValueError( + f"dflash_layer_ids {sorted(missing)} not reached -- " + f"model only has {len(self.layers)} layers" + ) + hidden = torch.cat([captured[i] for i in self.dflash_layer_ids], dim=-1) + + x = self.norm(x) + logits = self.lm_head(x).float() # (B, T, V) -- NOT x[:, -1, :] + cap = self.logit_softcap.float() + logits = torch.tanh(logits / cap) * cap + return logits, hidden + + model.forward = types.MethodType(_mlx_dflash_model_forward, model) + + +def dflash_mlx_source_transformations( + model: nn.Module, + dtype: torch.dtype = torch.bfloat16, + use_turboquant: bool = False, + max_write_len: int | None = None, +) -> None: + """Apply MLX source transformations to a Gemma4_31BWithHidden model in-place. + + Identical to mlx_source_transformations.mlx_source_transformations() + except the final step installs a DFlash-aware top-level forward + (full-sequence logits + hidden-state capture) instead of the stock + last-token-only, single-output one. See module docstring for why this + duplication is necessary rather than composing with the original. + + model must be a Gemma4_31BWithHidden instance (needs .dflash_layer_ids). + """ + if not hasattr(model, "dflash_layer_ids"): + raise TypeError( + "dflash_mlx_source_transformations requires a model with " + "dflash_layer_ids (e.g. Gemma4_31BWithHidden), got " + f"{type(model).__name__}" + ) + + config = model.config + + for layer in model.layers: + attn = layer.self_attn + + if attn.is_sliding: + sliding_write_len = ( + min(max_write_len, config.sliding_window) + if max_write_len is not None + else None + ) + attn.kv_cache = MLXRingKVCache( + max_batch_size=1, + max_context_length=config.sliding_window, + n_heads=attn.n_kv_heads, + head_dim=attn.head_dim, + dtype=dtype, + max_write_len=sliding_write_len, + ) + attn.is_turboquant = False + elif use_turboquant: + attn.kv_cache = MLXTurboQuantKVCache( + max_batch_size=1, + max_context_length=config.max_seq_len, + n_heads=attn.n_kv_heads, + head_dim=attn.head_dim, + enable_dynamic_shape=True, + dtype=dtype, + ) + attn.is_turboquant = True + else: + attn.kv_cache = MLXKVCache( + max_batch_size=1, + max_context_length=config.max_seq_len, + n_heads=attn.n_kv_heads, + head_dim=attn.head_dim, + enable_dynamic_shape=True, + dtype=dtype, + ) + attn.is_turboquant = False + + _replace_attention_forward(attn) + _replace_layer_forward(layer) + + _replace_dflash_model_forward(model) diff --git a/examples/models/gemma4_31b/export_dflash_draft.py b/examples/models/gemma4_31b/export_dflash_draft.py new file mode 100644 index 00000000000..61d004bca32 --- /dev/null +++ b/examples/models/gemma4_31b/export_dflash_draft.py @@ -0,0 +1,185 @@ +# 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 Gemma4-31B DFlash draft model to a .pte program. + +Mirrors examples/models/qwen3/export_dflash_draft.py's flow (DFlashDraftModel +and load_dflash_config are fully generic -- no Gemma4-specific changes +needed there), but pulls the shared embed_tokens/lm_head weights from the +*prequantized* target checkpoint instead of loading the full bf16 target +model via AutoModelForCausalLM. For a 31B model, loading the full bf16 +target (~62GB resident) just to copy two weight tensors before discarding +it would resurrect the exact memory problem the prequantized-checkpoint +export path was chosen to avoid. + +embed_tokens and lm_head are stored as separately-quantized tensors in the +prequantized checkpoint (confirmed via direct safetensors key inspection: +embed_tokens._weight_qdata/_scale/_zero_point and lm_head._weight_qdata/... +both present) -- config.json's tie_word_embeddings: true is stale/does not +reflect this; do not rely on it. Each is unflattened back into its torchao +tensor subclass via unflatten_tensor_state_dict (same utility +quant/pack_mlx.py's load_and_pack_for_mlx uses) and dequantized to bf16 +before being copied into the draft model's own embed_tokens/lm_head. +""" + +import argparse +import json + +import torch + +from executorch.backends.mlx.examples.llm.dflash_draft_model import ( + DFlashDraftModel, + load_dflash_config, +) +from executorch.examples.models.gemma4_31b.quant.quantize import dequantize_weight +from huggingface_hub import snapshot_download +from pathlib import Path +from safetensors import safe_open +from safetensors.torch import load_file +from torch.export import Dim + + +def _load_dequantized_tensor( + safetensors_path: str, logical_name: str, dtype: torch.dtype = torch.bfloat16 +) -> torch.Tensor: + """Load one named weight from a torchao-quantized safetensors checkpoint, + reconstructing the tensor subclass then dequantizing to a dense tensor. + + logical_name is a dotted module path, e.g. "embed_tokens.weight" or + "lm_head.weight" -- matches the tensor_names entries in the checkpoint's + safetensors metadata, not the raw on-disk key prefixes. + """ + from torchao.prototype.safetensors.safetensors_support import ( + unflatten_tensor_state_dict, + ) + + with safe_open(safetensors_path, framework="pt", device="cpu") as f: + metadata = f.metadata() + all_keys = list(f.keys()) + tensor_names = json.loads(metadata.get("tensor_names", "[]")) + if logical_name not in tensor_names: + raise KeyError( + f"{logical_name!r} not found in checkpoint tensor_names " + f"(have: {[n for n in tensor_names if 'embed' in n or 'lm_head' in n]})" + ) + parts = logical_name.rsplit(".", 1) + module_fqn, weight_name = parts[0], parts[-1] + prefix = f"{module_fqn}._{weight_name}_" + partial = {k: f.get_tensor(k) for k in all_keys if k.startswith(prefix)} + if not partial: + raise KeyError( + f"No keys found with prefix {prefix!r} for {logical_name!r}" + ) + result, _ = unflatten_tensor_state_dict(partial, metadata) + + reconstructed = result[logical_name] + return dequantize_weight(reconstructed, dtype=dtype) + + +def load_draft_model_from_prequantized_target( + draft_id: str, target_prequantized_dir: str +) -> 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}" + + target_safetensors = f"{target_prequantized_dir}/model.safetensors" + print(f"Dequantizing embed_tokens.weight from {target_safetensors}...") + embed_weight = _load_dequantized_tensor(target_safetensors, "embed_tokens.weight") + model.embed_tokens.weight.data.copy_(embed_weight) + + print(f"Dequantizing lm_head.weight from {target_safetensors}...") + lm_head_weight = _load_dequantized_tensor(target_safetensors, "lm_head.weight") + model.lm_head.weight.data.copy_(lm_head_weight) + + return model + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--target-prequantized", + default="./gemma-4-31B-it-HQQ-INT4", + help="Directory with the prequantized target checkpoint " + "(model.safetensors + config.json) -- used only to source " + "embed_tokens/lm_head weights, not to build the full target model.", + ) + parser.add_argument("--draft-model", default="z-lab/gemma-4-31B-it-DFlash") + parser.add_argument("--output", default="gemma4_31b_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() + + model = load_draft_model_from_prequantized_target( + args.draft_model, args.target_prequantized + ) + model.eval() + + # Quantize the draft model to match the target model's precision -- + # keeping both at the same precision reduces memory and helps keep + # their predictions consistent, important for a high 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/gemma4_31b/export_dflash_target.py b/examples/models/gemma4_31b/export_dflash_target.py new file mode 100644 index 00000000000..d53fa77224e --- /dev/null +++ b/examples/models/gemma4_31b/export_dflash_target.py @@ -0,0 +1,211 @@ + +"""Export Gemma4-31B as a DFlash target model (logits + hidden states) to +the MLX backend. + +Mirrors export.py's load_prequantized_model + _export_mlx flow, but builds +Gemma4_31BWithHidden instead of Gemma4_31B, and skips the sample=True / +SamplingHead branch entirely -- DFlash verification needs raw per-position +logits (see dflash_hidden_export.py's forward docstring for why), not an +on-device sampled token. + +layer_ids should match the DFlash draft checkpoint's config.json +dflash_config.target_layer_ids exactly (e.g. z-lab/gemma-4-31B-it-DFlash: +[1, 12, 23, 35, 46, 57]) -- a mismatch here would silently condition the +draft on the wrong hidden states with no error, only degraded tau. + +Usage: + python3 export_dflash_target.py \\ + --prequantized ./gemma-4-31B-it-HQQ-INT4 \\ + --dflash-layers 1,12,23,35,46,57 \\ + --output-dir ./gemma4_31b_dflash_exports_mlx \\ + --max-seq-len 4096 +""" + +import argparse +import gc +import json +import os + +import torch + +from executorch.examples.models.gemma4_31b.dflash_hidden_export import ( + Gemma4_31BWithHidden, +) +from executorch.examples.models.gemma4_31b.export import _pack_for_backend +from executorch.examples.models.gemma4_31b.model import ( + Gemma4_31BConfig, + materialize_runtime_buffers, +) + + +def load_prequantized_dflash_target( + prequantized_dir: str, + layer_ids: list, + max_seq_len: int = 4096, +) -> tuple: + """Load a quantized checkpoint into Gemma4_31BWithHidden, packed for MLX. + + Same as export.py's load_prequantized_model, except the meta-device + model is Gemma4_31BWithHidden (adds no new parameters/buffers beyond + Gemma4_31B, so _pack_for_backend's generic by-FQN packing is unaffected). + """ + config = Gemma4_31BConfig.from_hf_config( + os.path.join(prequantized_dir, "config.json") + ) + config.max_seq_len = max_seq_len + + print("Building Gemma4_31BWithHidden on meta device...") + with torch.device("meta"): + model = Gemma4_31BWithHidden(config, layer_ids=layer_ids) + + safetensors_path = os.path.join(prequantized_dir, "model.safetensors") + print(f"Loading quantized checkpoint from {safetensors_path}...") + _pack_for_backend(model, safetensors_path, "mlx") + model.eval() + + print( + f"Model: {config.num_hidden_layers} layers, hidden={config.hidden_size}, " + f"dflash_layer_ids={layer_ids}" + ) + return model, config + + +def export_dflash_target_mlx( + model: "Gemma4_31BWithHidden", + config: Gemma4_31BConfig, + output_dir: str, +) -> None: + """Export the DFlash target (logits, hidden) via torch.export + MLX backend. + + Adapted from export.py's _export_mlx: same source transforms and + lowering pipeline, but no sample=True / SamplingHead branch (not + applicable -- see load call site), and example_args/dynamic_shapes + reflect forward(tokens, input_pos) -> (logits, hidden) instead of + forward(tokens, input_pos, temperature) -> sampled_token. + """ + import executorch.backends.mlx.custom_kernel_ops.gguf.patterns # noqa: F401 + import executorch.extension.llm.export.gguf # noqa: F401 + import executorch.extension.llm.export.int4 # noqa: F401 + + from executorch.backends.mlx import MLXPartitioner + from executorch.backends.mlx.passes import get_default_passes + from executorch.examples.models.gemma4_31b.dflash_mlx_source_transformations import ( + dflash_mlx_source_transformations, + ) + from executorch.exir import ( + EdgeCompileConfig, + ExecutorchBackendConfig, + to_edge_transform_and_lower, + ) + from executorch.exir.passes import MemoryPlanningPass + from torch.export import Dim, export + + # block_size (max verified block length) governs the upper bound here, + # not max_prefill in the general sense -- DFlash target verification + # forward passes are always <= block_size tokens. 256 matches export.py's + # ceiling and comfortably covers any realistic block_size (e.g. 16). + max_verify_len = 256 + + dflash_mlx_source_transformations( + model, + dtype=torch.bfloat16, + use_turboquant=False, + max_write_len=max_verify_len, + ) + + materialize_runtime_buffers(model, dtype=torch.bfloat16) + + seq_dim = Dim("seq_len", min=1, max=max_verify_len) + example_tokens = torch.tensor([[0, 1]], dtype=torch.long) + example_input_pos = torch.tensor([0, 1], dtype=torch.long) + example_args = (example_tokens, example_input_pos) + dynamic_shapes = ({1: seq_dim}, {0: seq_dim}) + + print(f"Exporting DFlash target (T in [1, {max_verify_len}])...") + with torch.no_grad(): + exported = export( + model, + example_args, + dynamic_shapes=dynamic_shapes, + strict=True, + ) + + del model + gc.collect() + + print("Lowering to ExecuTorch with MLX backend...") + et_prog = to_edge_transform_and_lower( + exported, + transform_passes=get_default_passes(), + partitioner=[MLXPartitioner()], + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ), + constant_methods={ + "get_max_seq_len": config.max_seq_len, + "get_vocab_size": config.vocab_size, + "get_n_layers": config.num_hidden_layers, + "get_max_prefill_chunk": max_verify_len, + "use_kv_cache": True, + "use_sdpa_with_kv_cache": False, + "enable_dynamic_shape": True, + "use_sampling": False, + }, + ) + + del exported + gc.collect() + et_program = et_prog.to_executorch( + config=ExecutorchBackendConfig( + extract_delegate_segments=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + ), + ) + + del et_prog + gc.collect() + + os.makedirs(output_dir, exist_ok=True) + pte_path = os.path.join(output_dir, "model.pte") + print(f"Saving to {pte_path}...") + with open(pte_path, "wb") as f: + et_program.write_to_file(f) + print(f" {os.path.getsize(pte_path) / 1024**2:.1f} MB") + if et_program._tensor_data: + et_program.write_tensor_data_to_file(output_dir) + print(f" Saved tensor data (.ptd) to {output_dir}/") + print("Done.") + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--prequantized", + required=True, + help="Directory with quantized checkpoint (model.safetensors + config.json).", + ) + p.add_argument( + "--dflash-layers", + required=True, + help=( + "Comma-separated 0-indexed target layer ids to tap for hidden " + "states, matching the DFlash draft checkpoint's config.json " + "dflash_config.target_layer_ids exactly (e.g. " + "'1,12,23,35,46,57' for z-lab/gemma-4-31B-it-DFlash)." + ), + ) + p.add_argument("--output-dir", required=True, help="Output dir for model.pte/.ptd.") + p.add_argument("--max-seq-len", type=int, default=4096) + args = p.parse_args() + + layer_ids = [int(x) for x in args.dflash_layers.split(",")] + + model, config = load_prequantized_dflash_target( + args.prequantized, layer_ids, max_seq_len=args.max_seq_len + ) + export_dflash_target_mlx(model, config, args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/examples/models/gemma4_31b/run_baseline.py b/examples/models/gemma4_31b/run_baseline.py new file mode 100644 index 00000000000..23cde0c3bf6 --- /dev/null +++ b/examples/models/gemma4_31b/run_baseline.py @@ -0,0 +1,102 @@ +# 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. + +Mirrors examples/models/qwen3/run_baseline.py's structure (same target .pte +reused for both baseline and DFlash -- see run_dflash.py's module docstring +for why), but Gemma4-31B's tokenizer/chat-template handling genuinely +differs from Qwen3's and is NOT a drop-in AutoTokenizer.apply_chat_template +call: + +- Uses the raw `tokenizers.Tokenizer` API (`Tokenizer.from_file` + + `.encode(...).ids` / `.decode(...)`), not transformers.AutoTokenizer -- + this checkpoint's tokenizer_config.json has no chat_template field at all. +- BOS is not part of the chat-template string; it must be prepended at the + token-ID level or, per inference.py's docstring, "the model's logits + collapse to a single high-frequency vocab token." +- Three stop tokens, not one: eos_token_ids = {1, 50, 106} (hardcoded in + inference.py's main(), not derived from the tokenizer's own eos_token + field). + +apply_chat_template is imported directly from inference.py rather than +reimplemented, so this only has one place to go stale. +""" + +import argparse +import time + +import torch +from executorch.examples.models.gemma4_31b.inference import apply_chat_template +from executorch.runtime import Runtime, Verification +from tokenizers import Tokenizer + +EOS_TOKEN_IDS = {1, 50, 106} +BOS_TOKEN_ID = 2 + + +def main(): + p = argparse.ArgumentParser() + p.add_argument( + "--target-pte", default="gemma4_31b_dflash_exports_mlx/model.pte" + ) + p.add_argument( + "--tokenizer-path", + default="./gemma-4-31B-it-HQQ-INT4/tokenizer.json", + ) + 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, + ) + args = p.parse_args() + + tokenizer = Tokenizer.from_file(args.tokenizer_path) + + prompt_str = args.prompt if not args.chat_template else apply_chat_template(args.prompt) + input_ids = tokenizer.encode(prompt_str).ids + if not input_ids or input_ids[0] != BOS_TOKEN_ID: + input_ids = [BOS_TOKEN_ID] + input_ids + prompt_ids = torch.tensor([input_ids], dtype=torch.long) + + 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 in EOS_TOKEN_IDS: + 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/gemma4_31b/run_dflash.py b/examples/models/gemma4_31b/run_dflash.py new file mode 100644 index 00000000000..44ab0e2f77d --- /dev/null +++ b/examples/models/gemma4_31b/run_dflash.py @@ -0,0 +1,188 @@ +# 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 Gemma4-31B. + +Mirrors examples/models/qwen3/run_dflash.py's algorithm exactly (see that +file's module docstring for the four-step round structure and V1 scope +notes -- unchanged here). Tokenizer/chat-template handling differs +genuinely from Qwen3 -- see run_baseline.py's module docstring for why +(raw tokenizers.Tokenizer API, manual BOS prepend, three hardcoded EOS ids, +apply_chat_template imported directly from inference.py). +""" + +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.examples.models.gemma4_31b.inference import apply_chat_template +from executorch.runtime import Runtime, Verification +from tokenizers import Tokenizer + +EOS_TOKEN_IDS = {1, 50, 106} +BOS_TOKEN_ID = 2 + + +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="gemma4_31b_dflash_exports_mlx/model.pte" + ) + p.add_argument("--draft-pte", default="gemma4_31b_dflash_draft.pte") + p.add_argument( + "--draft-config-dir", + default="./gemma-4-31B-it-DFlash", + help="Local directory with the draft checkpoint's config.json " + "(from `hf download z-lab/gemma-4-31B-it-DFlash --local-dir ...`).", + ) + p.add_argument( + "--tokenizer-path", + default="./gemma-4-31B-it-HQQ-INT4/tokenizer.json", + ) + 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 Gemma4's chat template. On by default (matches the " + "draft model's training distribution).", + ) + 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.", + ) + args = p.parse_args() + + config = load_dflash_config(Path(args.draft_config_dir)) + mask_id = config.mask_token_id + block_size = args.block_size if args.block_size is not None else config.block_size + + tokenizer = Tokenizer.from_file(args.tokenizer_path) + + 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") + + prompt_str = args.prompt if not args.chat_template else apply_chat_template(args.prompt) + input_ids = tokenizer.encode(prompt_str).ids + if not input_ids or input_ids[0] != BOS_TOKEN_ID: + input_ids = [BOS_TOKEN_ID] + input_ids + prompt_ids = torch.tensor([input_ids], dtype=torch.long) + prompt_len = prompt_ids.shape[1] + + 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 + bs = block_size + + 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() + + 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() + + 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]] + + hit_eos = [t for t in new_tokens if t in EOS_TOKEN_IDS] + if hit_eos: + first_eos_pos = min(new_tokens.index(t) for t in hit_eos) + new_tokens = new_tokens[: first_eos_pos + 1] + accepted = min(accepted, len(new_tokens) - 1) + + accepted_total += accepted + emitted_total += len(new_tokens) + + generated.extend(new_tokens) + + pos += len(new_tokens) + last_token = new_tokens[-1] + hidden = torch.cat([hidden, new_hidden[:, : len(new_tokens), :].float()], dim=1) + + if last_token in EOS_TOKEN_IDS: + 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/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..bae527766fc --- /dev/null +++ b/examples/models/qwen3/export_dflash_draft.py @@ -0,0 +1,120 @@ +# 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) + # --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) + 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/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..451ebbaef17 --- /dev/null +++ b/examples/models/qwen3/run_dflash.py @@ -0,0 +1,221 @@ +# 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]] + + # 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}") + 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)}")