# How DFlash Integrates with Qwen3 Rotary Embeddings for Speculative Decoding > Discover how DFlash integrates Qwen3 rotary embeddings using the Hugging Face transformers library. Learn how to ensure exact positional encoding matches for optimized draft models. - Repository: [Z Lab/dflash](https://github.com/z-lab/dflash) - Tags: how-to-guide - Published: 2026-04-17 --- **DFlash imports the `Qwen3RotaryEmbedding` class directly from the Hugging Face transformers library and instantiates it within `DFlashDraftModel` to ensure the draft model's positional encoding matches the target Qwen3 model exactly.** The `z-lab/dflash` repository implements a speculative decoding framework that relies on a lightweight draft model to predict tokens ahead of the main Qwen3 target model. To maintain the architectural consistency required for speculative decoding, DFlash integrates Qwen3 rotary embeddings by reusing the native implementation rather than creating a custom variant. ## Reusing Qwen3RotaryEmbedding in the Draft Model DFlash builds its draft model on top of the Qwen3 architecture by importing the positional encoding components directly from the transformers library. The `Qwen3RotaryEmbedding` class is imported from `transformers.models.qwen3.modeling_qwen3` and instantiated during model construction. When a `DFlashDraftModel` is created, it initializes the rotary embedding object using the model configuration: ```python # dflash/model.py, line 316 self.rotary_emb = Qwen3RotaryEmbedding(config) ``` This instantiation ensures the draft model uses the same rotary embedding dimensions, base frequencies, and scaling factors as the original Qwen3 architecture. ## Forward Pass and Position Embedding Generation During the forward pass, the draft model generates positional embeddings by invoking the rotary embedding layer with the current hidden states and position IDs supplied by the caller. This occurs in the `forward` method of `DFlashDraftModel`: ```python # dflash/model.py, line 335 position_embeddings = self.rotary_emb(hidden_states, position_ids) ``` The call returns a tuple `(cos, sin)` containing pre-computed cosine and sine tensors for the given sequence length. These tensors represent the rotary positional encoding values that will be applied to the query and key vectors in the attention mechanism. ## Applying Embeddings in Attention Layers ### The Qwen3DFlashAttention Module Each decoder layer in the draft model receives `position_embeddings` and passes them to the custom `Qwen3DFlashAttention` module. This module implements the attention mechanism specific to the DFlash speculative decoding framework while maintaining compatibility with Qwen3's rotary encoding scheme. ### Rotary Position Embedding Application Inside the attention module, the cosine and sine tensors are applied to the query and key tensors via the helper function `apply_rotary_pos_emb`. This function rotates the query and key vectors by the angles specified in the positional encoding: ```python # dflash/model.py, lines 34-38 (apply_rotary_pos_emb) q, k = apply_rotary_pos_emb(q, k, cos, sin) ``` This implementation mirrors the exact rotary-positional encoding logic used by the original Qwen3 model, ensuring that the draft model's attention operates on correctly rotated token representations. ## Alignment with the Target Model The `Qwen3RotaryEmbedding` class is identical to the one used by the target (full) Qwen3 model, ensuring both draft and target share the same positional encodings. This alignment is essential for the speculative decoding algorithm employed by DFlash, as any discrepancy in positional encoding would cause the draft model's predictions to diverge from the target model's probability distribution. ## Complete Integration Example The following example demonstrates how DFlash instantiates the draft model and utilizes the rotary embedding during inference: ```python import torch from transformers import AutoConfig from dflash.model import DFlashDraftModel # ------------------------------------------------- # 1️⃣ Load a Qwen‑3 config (e.g., qwen3‑4b‑instruct) # ------------------------------------------------- config = AutoConfig.from_pretrained("Qwen/Qwen3-4B-Instruct") # Enable DFlash‑specific settings (example values) config.dflash_config = { "target_layer_ids": [2, 6, 10], # layers whose hidden states are extracted "mask_token_id": 0, "block_size": 4, } config.num_target_layers = 12 # number of layers in the full model config.block_size = 4 # speculative block size # ------------------------------------------------- # 2️⃣ Build the DFlash draft model # ------------------------------------------------- draft = DFlashDraftModel(config).eval().to("cuda") # ------------------------------------------------- # 3️⃣ Prepare dummy inputs # ------------------------------------------------- batch_size = 1 seq_len = 8 input_ids = torch.arange(seq_len).unsqueeze(0).to("cuda") position_ids = torch.arange(seq_len).unsqueeze(0).to("cuda") # DFlash requires the *noise embedding* (tokens from the target model) and # the *target hidden* representation. For illustration we use random tensors. noise_embedding = torch.randn(batch_size, seq_len, config.hidden_size, device="cuda") target_hidden = torch.randn( batch_size, len(config.dflash_config["target_layer_ids"]), config.hidden_size, device="cuda" ) # ------------------------------------------------- # 4️⃣ Forward pass – note the rotary embedding usage # ------------------------------------------------- with torch.no_grad(): logits = draft( position_ids=position_ids, attention_mask=None, noise_embedding=noise_embedding, target_hidden=target_hidden, ) print("logits shape:", logits.shape) # → (1, seq_len, vocab_size) ``` In this example, `self.rotary_emb = Qwen3RotaryEmbedding(config)` creates the positional encoder, and `self.rotary_emb(hidden_states, position_ids)` generates the `(cos, sin)` tensors applied to attention layers. ## Summary - DFlash imports `Qwen3RotaryEmbedding` directly from `transformers.models.qwen3.modeling_qwen3` to avoid implementation drift. - The rotary embedding is instantiated in `DFlashDraftModel` at [`dflash/model.py`](https://github.com/z-lab/dflash/blob/main/dflash/model.py) line 316 and invoked during the forward pass at line 335. - The embedding generates `(cos, sin)` tensors that are passed to `Qwen3DFlashAttention` layers. - The `apply_rotary_pos_emb` function applies these tensors to query and key vectors at lines 34-38 of [`dflash/model.py`](https://github.com/z-lab/dflash/blob/main/dflash/model.py). - Using the identical rotary embedding implementation ensures architectural consistency between the draft and target models for accurate speculative decoding. ## Frequently Asked Questions ### Why does DFlash use the native Qwen3RotaryEmbedding instead of creating a custom implementation? DFlash uses the native `Qwen3RotaryEmbedding` class to ensure perfect alignment with the target Qwen3 model's positional encoding scheme. Since speculative decoding requires the draft model's probability distribution to match the target model's distribution for valid token acceptance, any deviation in rotary embedding implementation would cause prediction drift and reduce decoding efficiency. ### How are the cosine and sine tensors generated by rotary embeddings used during attention computation? The `Qwen3RotaryEmbedding` returns a tuple of `(cos, sin)` tensors that represent the rotational angles for each position in the sequence. Inside the `Qwen3DFlashAttention` module, these tensors are passed to the `apply_rotary_pos_emb` function, which rotates the query and key vectors by these angles before computing the attention scores, thereby encoding positional information directly into the attention mechanism. ### What happens if the draft and target models use different rotary embedding implementations? If the draft and target models use different rotary embedding implementations, the hidden states and probability distributions between the two models will diverge. In speculative decoding, this causes the draft model's predicted tokens to be rejected by the target model at a higher rate, defeating the purpose of the speculative mechanism and potentially making inference slower than standard autoregressive generation. ### Does DFlash modify the Qwen3RotaryEmbedding class or use it as-is? DFlash uses the `Qwen3RotaryEmbedding` class as-is without modification. The class is imported directly from `transformers.models.qwen3.modeling_qwen3` and instantiated in `DFlashDraftModel` using the standard configuration object. This unmodified reuse ensures that the rotary embedding behavior remains identical to the upstream Qwen3 implementation.