z-lab / dflash

z-lab/dflash

open_in_new前往仓库

DFlash：用于闪存推测解码的块扩散方法

项目概览

项目概述

DFlash 是一个轻量级的块扩散（Block Diffusion）模型，专为推测解码（Speculative Decoding）设计，旨在实现高效、高质量的并行草稿生成。它通过引入块扩散机制，能够同时生成多个 token 的草稿，从而加速大语言模型的推理过程。项目提供了与多种主流推理后端（如 vLLM、SGLang、Transformers、MLX）的集成，并支持多种主流模型（如 Qwen、Gemma、MiniMax、Kimi 等）。

核心功能

  • 块扩散推测解码：采用块扩散模型作为草稿模型，实现并行 token 生成，显著提升解码速度。
  • 多后端支持：支持 vLLM、SGLang、Transformers 和 MLX（Apple Silicon）等推理后端。
  • 广泛模型兼容：提供针对多种主流模型的预训练 DFlash 草稿模型，并支持用户请求新增模型。
  • 简单易用的 API：提供 Python 接口和命令行工具，方便集成和基准测试。
  • 评估基准：内置基准测试脚本，支持 gsm8k、math500、humaneval 等数据集。

适用与不适用场景

适用场景：

  • 需要加速大语言模型推理的在线服务或离线批处理场景。
  • 使用支持列表中的模型（如 Qwen、Gemma、MiniMax、Kimi 等）进行部署。
  • 希望在不牺牲太多生成质量的前提下提升吞吐量的场景。

不适用场景：

  • 需要支持列表中未列出的模型，且用户无法自行训练 DFlash 草稿模型。
  • 对推理延迟极度敏感且无法接受额外草稿模型内存开销的场景。
  • 使用不支持的后端（如 TensorRT-LLM 等）或硬件平台（除 Apple Silicon 外）。

技术架构与依赖

DFlash 基于块扩散模型架构，通过扩散过程生成 token 块。其核心依赖包括：

  • Python 3.10+（推测）
  • PyTorch（基础框架）
  • Transformers（用于 Transformers 后端）
  • vLLM（用于 vLLM 后端，需 v0.20.1+）
  • SGLang（用于 SGLang 后端）
  • MLX（用于 Apple Silicon 后端）
  • uv（推荐包管理器）

具体依赖版本和安装方式请参考官方文档。

安装与快速开始

安装

推荐使用虚拟环境，根据后端选择安装命令：

  • Transformers: uv pip install -e ".[transformers]"
  • SGLang: uv pip install -e ".[sglang]"
  • vLLM: uv pip install -e ".[vllm]"（vLLM v0.20.1+ 已内置 DFlash 支持）
  • MLX: pip install -e ".[mlx]"

对于 Gemma4 模型，可能需要使用临时 Docker 镜像或源码安装。

快速开始

以 vLLM 为例，启动服务：

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

其他后端的快速开始示例请参考官方 README。

典型使用方法

Transformers 后端

from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer

draft = AutoModel.from_pretrained("z-lab/Qwen3-8B-DFlash-b16", trust_remote_code=True, dtype="auto", device_map="cuda:0").eval()
target = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-8B", dtype="auto", device_map="cuda:0").eval()
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")

messages = [{"role": "user", "content": "How many positive whole-number divisors does 196 have?"}]
input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True, enable_thinking=False).to(draft.device)

output = draft.spec_generate(input_ids=input_ids, max_new_tokens=2048, temperature=0.0, target=target, stop_token_ids=[tokenizer.eos_token_id])
print(tokenizer.decode(output[0], skip_special_tokens=False))

MLX 后端

from dflash.model_mlx import load, load_draft, stream_generate

model, tokenizer = load("Qwen/Qwen3.5-4B")
draft = load_draft("z-lab/Qwen3.5-4B-DFlash")

messages = [{"role": "user", "content": "How many positive whole-number divisors does 196 have?"}]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=True)
for r in stream_generate(model, draft, tokenizer, prompt, block_size=16, max_tokens=2048, temperature=0.6):
    print(r.text, end="", flush=True)

配置与部署要点

  • vLLM 配置：通过 --speculative-config 指定 DFlash 方法、草稿模型和推测 token 数。Gemma4 模型可能需要特殊 Docker 镜像。
  • SGLang 配置：设置环境变量 SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1，并通过启动参数指定推测算法和草稿模型。
  • Transformers 后端：仅支持 Qwen3 和 LLaMA-3.1 模型，需使用 trust_remote_code=True。
  • MLX 后端：适用于 Apple Silicon，需安装 MLX 依赖。
  • 部署建议：建议使用 GPU 环境，并确保显存充足；对于大规模部署，可考虑使用 Docker 容器。

限制、风险与许可证

  • 限制：目前仅支持特定模型列表；Transformers 后端仅支持部分模型；部分后端（如 SGLang 的 schedule overlapping）仍处于实验阶段，可能不稳定。
  • 风险：推测解码可能引入额外的内存开销和推理延迟；草稿模型质量可能影响生成质量。
  • 许可证：项目采用 MIT 许可证，允许自由使用、修改和分发。

官方链接

信息来源和分析时间