DFlash：块扩散模型加速 LLM 推理

class BlockDiffusionDraft:
    """块扩散 draft 模型"""

    def __init__(self, draft_model, block_size=16):
        self.draft = draft_model
        self.block_size = block_size

    def draft_tokens(self, context: Tensor) -> list[str]:
        # 1. 对目标位置加噪
        noise_level = sample_noise_levels(self.block_size)
        noisy_tokens = add_noise(context, noise_level)

        # 2. 块级别去噪——一次前向，不是自回归
        denoised = self.draft(noisy_tokens, context)

        # 3. 采样得到 token 序列
        draft_tokens = sample_tokens(denoised, temperature=0.0)
        return draft_tokens

# Transformers 后端
uv pip install -e ".[transformers]"

# SGLang 后端
uv pip install -e ".[sglang]"

# vLLM 后端（需要 nightly 版本）
uv pip install -e ".[vllm]"
uv pip install -U vllm --torch-backend=auto --extra-index-url https://wheels.vllm.ai/nightly

# Apple Silicon (MLX)
pip install -e ".[mlx]"

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

export SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1

python -m sglang.launch_server \
    --model-path Qwen/Qwen3.5-35B-A3B \
    --speculative-algorithm DFLASH \
    --speculative-draft-model-path z-lab/Qwen3.5-35B-A3B-DFlash \
    --speculative-num-draft-tokens 16 \
    --tp-size 1 \
    --attention-backend trtllm_mha \
    --speculative-draft-attention-backend fa4 \
    --mem-fraction-static 0.75 \
    --mamba-scheduler-strategy extra_buffer \
    --trust-remote-code

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": "196 有多少个正约数？"}]
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))

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": "196 有多少个正约数？"}]
prompt = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True,
    enable_thinking=True
)

tps = 0.0
for r in stream_generate(
    model, draft, tokenizer, prompt,
    block_size=16, max_tokens=2048, temperature=0.6
):
    print(r.text, end="", flush=True)
    tps = r.generation_tps

print(f"\nThroughput: {tps:.2f} tok/s")

# vLLM 后端 benchmark
python -m dflash.benchmark \
    --backend vllm --base-url http://127.0.0.1:8000 \
    --model Qwen/Qwen3.5-27B --dataset gsm8k \
    --num-prompts 128 --concurrency 1 --enable-thinking

# SGLang 后端 benchmark
python -m dflash.benchmark \
    --backend sglang --base-url http://127.0.0.1:30000 \
    --model Qwen/Qwen3.5-35B-A3B --dataset gsm8k \
    --num-prompts 128 --concurrency 1 --enable-thinking

# Transformers 后端 benchmark
torchrun --nproc_per_node=8 -m dflash.benchmark \
    --backend transformers --model Qwen/Qwen3-8B \
    --draft-model z-lab/Qwen3-8B-DFlash-b16 \
    --dataset gsm8k --max-samples 128

def verify(draft_tokens, target_logits, temperature=0.0):
    target_probs = F.softmax(target_logits, dim=-1)
    if temperature == 0:
        # 贪婪解码：接受概率最高的 token
        accepted = torch.argmax(target_probs, dim=-1)
    else:
        # 采样：按概率分布随机采样
        accepted = torch.multinomial(target_probs, 1).squeeze(-1)
    return accepted[:len(draft_tokens)]

@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}
}