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]:
        """
        输入: context (已生成的token序列)
        输出: draft_tokens (预测的下一个block)
        """
        # 1. 加噪过程：模拟扩散的forward process
        noise_level = sample_noise_levels(self.block_size)
        noisy_tokens = add_noise(context, noise_level)
        
        # 2. 去噪过程：DFlash核心，一次前向生成多个token
        #    注意：这是"块级别"去噪，不是自回归
        denoised = self.draft(noisy_tokens, context)
        
        # 3. 采样得到token序列
        draft_tokens = sample_tokens(denoised, temperature=0.0)
        
        return draft_tokens  # 返回一个block的token

# 基础安装
uv pip install -e ".[transformers]"  # 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

# 可选：启用实验性特性
# 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-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模型
draft = AutoModel.from_pretrained(
    "z-lab/Qwen3-8B-DFlash-b16",
    trust_remote_code=True,
    dtype="auto",
    device_map="cuda:0"
).eval()

# 加载target模型
target = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B",
    dtype="auto",
    device_map="cuda:0"
).eval()

# 加载tokenizer
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)

# 使用DFlash生成
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")

# 构造prompt
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
)

# 流式生成
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")

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

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

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

class DFlashTrainer:
    """DFlash训练流程"""
    
    def training_step(self, batch):
        # 1. Target模型生成"正确"token序列
        target_output = self.target_model(batch.input_ids)
        target_tokens = target_output.token_ids
        
        # 2. 为Draft模型生成训练数据
        #    模拟"加噪-去噪"过程
        noisy_tokens = self.add_noise(target_tokens)
        
        # 3. Draft模型学习去噪
        draft_output = self.draft_model(noisy_tokens, batch.input_ids)
        
        # 4. 损失：让draft预测接近target
        loss = self.mse_loss(draft_output, target_tokens)
        
        # 5. 反向传播更新draft
        loss.backward()
        self.optimizer.step()
        
        return loss
    
    def add_noise(self, tokens):
        """模拟扩散的加噪过程"""
        # 关键：不是随机噪声，而是对token做"扰动"
        # 扰动方式：替换、删除、插入
        noise_level = torch.rand(len(tokens)) * self.max_noise
        return self.token_noiser(tokens, noise_level)

def verify(draft_tokens, target_logits, temperature=0.0):
    """
    验证draft tokens是否被接受
    """
    # 1. 计算target模型对每个draft token的概率
    target_probs = F.softmax(target_logits, dim=-1)
    
    # 2. 计算接受概率
    #    策略：贪婪（temperature=0）或采样（temperature>0）
    if temperature == 0:
        # 贪婪：直接接受概率最高的token
        accepted = torch.argmax(target_probs, dim=-1)
    else:
        # 采样：按照概率接受
        accepted = torch.multinomial(target_probs, 1).squeeze(-1)
    
    # 3. 返回接受的tokens和实际生成的tokens
    return accepted[:len(draft_tokens)]

# 1. 确认硬件支持
nvidia-smi  # 或 Apple Silicon: sysctl -n machdep.cpu.brand_string

# 2. 安装正确版本
pip install -e ".[vllm]"  # 或 sglang

# 3. 验证DFlash模型加载
python -c "from transformers import AutoModel; print('DFlash model loaded')"

# 4. 运行benchmark确认加速效果
python -m dflash.benchmark --backend vllm ...

# 5. 监控指标
#    - Token接受率 (target > 85%)
#    - 吞吐量提升 (target > 2x)
#    - 首token延迟 (应不增加)

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