它负责最终输出质量。DFlash 猜出来的 token 必须由它验证后才会被接受。
一张关系图
DFlash 不是替换 Gemma，也不是 attention backend。它是 speculative decoding 里的 drafter， 在运行时接入 Gemma 4 31B-it 的解码循环。
1
用户输入
prompt 进入 serving engine，主模型准备进行正常 autoregressive decode。
2
Gemma hidden states
Gemma 4 31B-it 提供 target hidden features，DFlash 会利用这些特征。
3
DFlash 猜一块
外置 DFlash speculator 并行预测 1 + num_speculative_tokens 个 query。
4
Gemma 验证
target model 并行验证候选 token，接受最长正确前缀。
5
循环加速
接受的 token 越多，单位 token 的 target forward 成本越低。
Gemma 4 31B-it 的结构重点
在 DFlash 语境里，Gemma 4 31B-it 是高质量的大模型 verifier。理解它的 attention 结构， 能解释为什么 MUSA PR 需要支持 head_dim 512、mm-prefix 和 non-causal FlashAttention。
Dense instruction model
31B-it 是 instruction-tuned dense target，不是 MoE。最终输出由它决定。
target model
verifier
Hybrid attention
大量 local sliding-window 层，间隔插入 full/global attention 层，兼顾长上下文和效率。
sliding window
global attention
Global head_dim 512
Gemma 4 31B 的 full/global attention 会触及 512 head dim，这影响 FlashAttention backend selection。
head_size <= 512
FLASH_ATTN
结构速记
Gemma 4 31B-it
role: target / verifier
shape: dense instruction-tuned model
attention: sliding-window layers + periodic full/global layers
dflash relation: provides hidden features and verifies DFlash guesses
important for MUSA: head_dim 512 + mm_prefix model selection + non-causal verify pathDFlash 到底是什么
它是单独的 speculator
DFlash 是外置 drafter/checkpoint，例如给 google/gemma-4-31B-it 配套的
gemma-4-31B-it-speculator.dflash。它不是集成在 Gemma 4 31B-it 权重里的模块。
它是 block-diffusion drafter
普通 speculative decoding 常见形式是小模型顺序猜多个 token；DFlash 则是用轻量 block diffusion drafter 并行猜一个 block。
普通自回归 decode
Gemma predicts token t
Gemma predicts token t+1
Gemma predicts token t+2
...质量稳定，但每个 token 都要走主模型，吞吐受限。
DFlash speculative decode
DFlash guesses a block
Gemma verifies the block
accept correct prefix
repeatDFlash 猜得越准，被 Gemma 一次接受的 token 越多，加速越明显。
为什么需要 non-causal attention
这是今天最重要的概念点：DFlash 的 drafter/verify block 不是普通单 token causal decode。 它一次处理一个并行 block，因此 attention backend 需要能处理 non-causal metadata。
Causal attention
第 i 个 query 只能看自己和左边，不能看右边。普通 decoder 生成下一个 token 时通常走这条路。
Non-causal block attention
DFlash block 里的多个 query 作为一个整体被处理，不能被限制成严格左到右的单 token decode。
代码层面的关键句
# DFlash cross-attention:
# context K/V from target hidden states
# Q from query embeddings (bonus + mask tokens)
num_query_per_req = 1 + num_speculative_tokens
common_attn_metadata.causal = False
因此，MUSA 的 FlashAttention backend 如果不声明并实现 supports_non_causal()，
vLLM 要么拒绝选择它，要么可能在错误 mask 下运行。
MooreThreads/vllm-musa PR #52 做了什么
PR #52 的核心不是“新增 DFlash 算法”，而是让 vllm-musa 能在 MUSA 上正确、高效地跑 Gemma 4 31B-it + DFlash。
| 改动 | 原因 | 学习重点 |
|---|---|---|
DFlashProposer.dummy_run 接受 common_attn_metadata |
MUSA runner 会把这个参数传给 drafter；原 dflash override 没接会 TypeError。 | runtime patch 要和 runner 调用签名保持一致。 |
| FULL CUDAGraph capture 默认开启 | DFlash draft loop 否则会被 kernel launch overhead 卡住。 | spec decode 不是只看算力，图捕获和 dispatch key 也很关键。 |
| capture sizes 改成 block-aligned | DFlash 每请求 query 是 1 + num_speculative_tokens，例如 8 个 spec token 时 block 为 9。 |
常见配置是 [9, 18, 36, 72]，对应 BS=1/2/4/8。 |
| MUSA FLASH_ATTN 支持 non-causal | DFlash verify/block attention 需要 causal=False。 |
backend capability 声明会直接影响 vLLM 的 backend selection。 |
| MUSA FLASH_ATTN 支持 head_size 512 | Gemma 4 31B 的 global/full attention 会用到 512 head dim。 | 模型结构会反向驱动 backend 能力边界。 |
处理 supports_mm_prefix=True |
Gemma 4 是 multimodal/mm-prefix 模型，文本场景也会影响 backend selection。 | 这里有残余风险：任意 partial 2D mask 并非完全覆盖。 |
容易混淆的概念
DFlash 不是 Gemma 自带模块
它是单独 speculator，运行时通过 serving framework 接入主模型。
MTP assistant 不是 DFlash
gemma-4-E2B-it-assistant 是 MTP drafter；DFlash 是 block-diffusion drafter。
Non-causal 不等于输出乱来
DFlash 猜测阶段可以 non-causal；最终接受仍由 Gemma 4 31B-it 验证。
复习清单
能说明 Gemma 4 31B-it 是 target/verifier，而 DFlash 是外置 drafter。
能解释 1 + num_speculative_tokens 为什么影响 CUDAGraph capture size。
能画出 causal attention 和 non-causal block attention 的差别。
能说出 PR #52 中 supports_non_causal=True 和 head_size <= 512 的动机。
知道 E2B-it assistant 是 MTP drafter，不是 DFlash。
知道 DFlash 的质量安全来自 target model verify，而不是 drafter 本身全权决定输出。
继续阅读
这些链接适合按顺序读，先看 PR，再看 vLLM DFlash 源码，最后回到模型卡。