Duplicate from z-lab/Qwen3.5-27B-DFlash

Co-authored-by: Jian Chen <jianchen0311@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/dflash_system.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,171 @@
+---
+license: mit
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- dflash
+- speculative-decoding
+- diffusion
+- efficiency
+- flash-decoding
+- qwen
+- diffusion-language-model
+---
+
+# Qwen3.5-27B-DFlash
+[**Paper**](https://arxiv.org/abs/2602.06036) | [**GitHub**](https://github.com/z-lab/dflash) | [**Blog**](https://z-lab.ai/projects/dflash/)
+
+**DFlash** is a novel speculative decoding method that utilizes a lightweight **block diffusion** model for drafting. It enables efficient, high-quality parallel drafting that pushes the limits of inference speed.
+
+This model is the **drafter** component. It must be used in conjunction with the target model `Qwen/Qwen3.5-27B`. It was trained with a context length of 4096 tokens.
+
+<div align="center">
+  <img src="assets/dflash_system.png" alt="DFlash Architecture" width="100%">
+</div>
+
+## Quick Start
+
+### Installation
+
+vLLM:
+```bash
+uv pip install vllm
+uv pip install -U vllm --torch-backend=auto --extra-index-url https://wheels.vllm.ai/nightly
+```
+
+SGLang:
+```bash
+uv pip install "git+https://github.com/sgl-project/sglang.git@refs/pull/20547/head#subdirectory=python"
+```
+
+### Launch Server
+
+vLLM:
+```bash
+vllm serve Qwen/Qwen3.5-27B \
+  --speculative-config '{"method": "dflash", "model": "z-lab/Qwen3.5-27B-DFlash", "num_speculative_tokens": 15}' \
+  --attention-backend flash_attn \
+  --max-num-batched-tokens 32768
+```
+
+SGLang:
+```bash
+# Optional: enable schedule overlapping (experimental, may not be stable)
+# export SGLANG_ENABLE_SPEC_V2=1
+# export SGLANG_ENABLE_DFLASH_SPEC_V2=1
+# export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1
+
+python -m sglang.launch_server \
+    --model-path Qwen/Qwen3.5-27B \
+    --speculative-algorithm DFLASH \
+    --speculative-draft-model-path z-lab/Qwen3.5-27B-DFlash \
+    --speculative-num-draft-tokens 16 \
+    --tp-size 1 \
+    --attention-backend fa3 \
+    --mem-fraction-static 0.75 \
+    --mamba-scheduler-strategy extra_buffer \
+    --trust-remote-code
+```
+> **Tip:** For long-context or agentic workloads, add `--speculative-dflash-draft-window-size WINDOW_SIZE` to enable sliding-window attention for the drafter.
+
+### Usage
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:30000/v1", api_key="EMPTY")
+
+response = client.chat.completions.create(
+    model="Qwen/Qwen3.5-27B",
+    messages=[{"role": "user", "content": "Write a quicksort in Python."}],
+    max_tokens=4096,
+    temperature=0.0
+)
+print(response.choices[0].message.content)
+```
+
+## Benchmark Results
+
+**Setup:** Single NVIDIA B200, SGLang, thinking enabled, max output length 4096. We report end-to-end throughput, including prefill time. See our [GitHub repository](https://github.com/z-lab/dflash) for reproduction scripts.
+
+### Throughput and Speedup
+
+_Tokens/sec (speedup vs. autoregressive baseline)_
+
+**Block Size = 16**
+
+| Task | Concurrency | AR | MTP | **DFlash** |
+|---|---:|---:|---:|---:|
+| Math500 | 1 | 84 | 243 (2.9x) | **397 (4.7x)** |
+| | 8 | 625 | 1457 (2.3x) | **2270 (3.6x)** |
+| | 16 | 1121 | 2224 (2.0x) | **3135 (2.8x)** |
+| | 32 | 1949 | 2504 (1.3x) | **3712 (1.9x)** |
+| GSM8K | 1 | 83 | 215 (2.6x) | **330 (4.0x)** |
+| | 8 | 625 | 1303 (2.1x) | **1868 (3.0x)** |
+| | 16 | 1109 | 1773 (1.6x) | **2589 (2.3x)** |
+| | 32 | 1914 | 2170 (1.1x) | **3152 (1.6x)** |
+| HumanEval | 1 | 83 | 236 (2.9x) | **427 (5.2x)** |
+| | 8 | 602 | 1345 (2.2x) | **2079 (3.5x)** |
+| | 16 | 1031 | 1921 (1.9x) | **2748 (2.7x)** |
+| | 32 | 1720 | 2234 (1.3x) | **3198 (1.9x)** |
+| MBPP | 1 | 84 | 200 (2.4x) | **347 (4.2x)** |
+| | 8 | 627 | 1049 (1.7x) | **1826 (2.9x)** |
+| | 16 | 1075 | 1729 (1.6x) | **2479 (2.3x)** |
+| | 32 | 1832 | 1933 (1.1x) | **2808 (1.5x)** |
+| MT-Bench | 1 | 84 | 169 (2.0x) | **255 (3.0x)** |
+| | 8 | 622 | 1035 (1.7x) | **1444 (2.3x)** |
+| | 16 | 1113 | 1550 (1.4x) | **1984 (1.8x)** |
+| | 32 | 1900 | 1772 (0.9x) | **2391 (1.3x)** |
+
+**Block Size = 8**
+
+| Task | Concurrency | AR | MTP | **DFlash** |
+|---|---:|---:|---:|---:|
+| Math500 | 1 | 84 | 273 (3.2x) | **335 (4.0x)** |
+| | 8 | 625 | 1673 (2.7x) | **2020 (3.2x)** |
+| | 16 | 1121 | 2731 (2.4x) | **3646 (3.3x)** |
+| | 32 | 1949 | 3739 (1.9x) | **4288 (2.2x)** |
+| GSM8K | 1 | 83 | 243 (2.9x) | **301 (3.6x)** |
+| | 8 | 625 | 1539 (2.5x) | **1814 (2.9x)** |
+| | 16 | 1109 | 2472 (2.2x) | **2896 (2.6x)** |
+| | 32 | 1914 | 3431 (1.8x) | **3822 (2.0x)** |
+| HumanEval | 1 | 83 | 258 (3.1x) | **350 (4.2x)** |
+| | 8 | 602 | 1486 (2.5x) | **1856 (3.1x)** |
+| | 16 | 1031 | 2302 (2.2x) | **2749 (2.7x)** |
+| | 32 | 1720 | 2477 (1.4x) | **3412 (2.0x)** |
+| MBPP | 1 | 84 | 234 (2.8x) | **311 (3.7x)** |
+| | 8 | 627 | 1375 (2.2x) | **1757 (2.8x)** |
+| | 16 | 1075 | 2159 (2.0x) | **2661 (2.5x)** |
+| | 32 | 1832 | 2885 (1.6x) | **3309 (1.8x)** |
+| MT-Bench | 1 | 84 | 210 (2.5x) | **250 (3.0x)** |
+| | 8 | 622 | 1300 (2.1x) | **1495 (2.4x)** |
+| | 16 | 1113 | 2105 (1.9x) | **2403 (2.2x)** |
+| | 32 | 1900 | 2873 (1.5x) | **3256 (1.7x)** |
+
+### Acceptance Length
+_Format: MTP / DFlash (averaged across concurrency levels)_
+
+| Task | B8 | B16 |
+|---|---:|---:|
+| Math500 | 5.73 / **5.90** | 7.14 / **7.93** |
+| GSM8K | 5.54 / **5.57** | 6.84 / **7.22** |
+| HumanEval | 5.81 / **6.34** | 7.38 / **9.18** |
+| MBPP | 5.10 / **5.60** | 5.94 / **7.27** |
+| MT-Bench | **4.60** / 4.54 | 5.30 / **5.47** |
+
+## Acknowledgements
+
+Special thanks to [David Wang](https://davidwa.ng/) for his outstanding engineering support on this project. We are also grateful to [Modal](https://modal.com/), [InnoMatrix](https://innomatrix.ai), and [Yotta Labs](https://www.yottalabs.ai/) for providing the compute resources used to train this draft model.
+
+## Citation
+
+If you find DFlash useful, please cite our work. To share feedback on DFlash or request new model support, please fill out this form: [DFlash Feedback](https://forms.gle/4YNwfqb4nJdqn6hq9).
+
+```bibtex
+@article{chen2026dflash,
+  title   = {{DFlash: Block Diffusion for Flash Speculative Decoding}},
+  author  = {Chen, Jian and Liang, Yesheng and Liu, Zhijian},
+  journal = {arXiv preprint arXiv:2602.06036},
+  year    = {2026}
+}
+```

config.json

@@ -0,0 +1,52 @@
+{
+  "architectures": [
+    "DFlashDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "dflash.DFlashDraftModel"
+  },
+  "block_size": 16,
+  "dflash_config": {
+    "mask_token_id": 248070,
+    "target_layer_ids": [
+      1,
+      16,
+      31,
+      46,
+      61
+    ]
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 248044,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 17408,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 262144,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 64,
+  "pad_token_id": 248044,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000000,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 248320
+}

dflash.py

@@ -0,0 +1,188 @@
+from typing import Optional, Callable
+from typing_extensions import Unpack, Tuple
+import torch
+from torch import nn
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+    Qwen3Config,
+    Qwen3PreTrainedModel,
+    Qwen3MLP,
+    GradientCheckpointingLayer,
+    FlashAttentionKwargs,
+    rotate_half,
+    eager_attention_forward,
+    ALL_ATTENTION_FUNCTIONS,
+)
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.cache_utils import Cache
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_len = q.size(-2)
+    q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :])
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+class Qwen3DFlashAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = False  
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+        self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        target_hidden: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        bsz, q_len = hidden_states.shape[:-1]
+        ctx_len = target_hidden.shape[1]
+        q = self.q_proj(hidden_states)
+        q = q.view(bsz, q_len, -1, self.head_dim)
+        q = self.q_norm(q).transpose(1, 2)
+        k_ctx = self.k_proj(target_hidden)
+        k_noise = self.k_proj(hidden_states)
+        v_ctx = self.v_proj(target_hidden)
+        v_noise = self.v_proj(hidden_states)
+        k = torch.cat([k_ctx, k_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim)
+        v = torch.cat([v_ctx, v_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim)
+        k = self.k_norm(k).transpose(1, 2)
+        v = v.transpose(1, 2)
+        cos, sin = position_embeddings
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+        if past_key_values is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs)
+        attn_fn: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        attn_output, attn_weights = attn_fn(
+            self,
+            q,
+            k,
+            v,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = Qwen3DFlashAttention(config=config, layer_idx=layer_idx)
+        self.mlp = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        target_hidden: Optional[torch.Tensor] = None,
+        hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            target_hidden=target_hidden,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )[0]
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+class DFlashDraftModel(Qwen3PreTrainedModel):
+    config_class = Qwen3Config
+    _no_split_modules = ["Qwen3DFlashDecoderLayer"]
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.config = config
+        self.layers = nn.ModuleList(
+            [Qwen3DFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.target_layer_ids = self.config.dflash_config.get("target_layer_ids", None)
+        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen3RotaryEmbedding(config)
+        self.fc = nn.Linear(len(self.target_layer_ids) * config.hidden_size, config.hidden_size, bias=False)
+        self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.block_size = config.block_size
+        self.mask_token_id = self.config.dflash_config.get("mask_token_id", None)
+        self.post_init()
+
+    def forward(
+        self,
+        position_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        noise_embedding: Optional[torch.Tensor] = None,
+        target_hidden: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        hidden_states = noise_embedding
+        target_hidden = self.hidden_norm(self.fc(target_hidden))
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states=hidden_states,
+                target_hidden=target_hidden,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                use_cache=use_cache,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+        return self.norm(hidden_states)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e44cae31ebda1940da56318c129509df468a1a6508a8f816a03e0c5c8661b77
+size 3460432504