o^t 和 self-conditioning 向量 S^t，prompt 侧 KV 在整个循环中保持静态。
DiffusionGemma 是 Google DeepMind 开源的离散扩散文本模型。它保留 Gemma 系列的语言主干，但把生成方式从逐 token 自回归改成了 256-token canvas 上的多步去噪。本章重点看三个问题：KV cache 还在不在，attention mask 到底是 causal 还是 full matrix，以及 256-token block 如何接成长文本。
传统自回归语言模型一次只生成一个 Token。Batch size 很小时，每一步都要从 HBM/VRAM 读一遍大块模型权重，却只对一个位置做计算，GPU 很容易卡在显存带宽上。
DiffusionGemma 一次处理一个固定长度的 canvas，默认是 256 个 Token。它不会在一次 forward 后就确定全部输出，而是在这个 canvas 上反复 denoise。
图像扩散可以给像素加高斯噪声。文本不行，Token "The" 没有“稍微变模糊”的连续状态，所以离散文本扩散通常会把部分 Token 替换成特殊符号或随机词表 Token。
类似于 BERT 的 Masked Language Modeling (MLM)，该机制在训练时将文本序列中的部分 Token 替换为特殊的 [MASK] 标记：
[MASK] 比例。[MASK] 位置的真实 Token。在每一步中，采样器根据模型预测的置信度，仅保留最确信的部分 Token，替换掉对应的 [MASK]，并在下一步继续预测剩余的 [MASK]。DiffusionGemma 使用均匀状态扩散：
[MASK] 标记，而是从词表中随机、均匀抽样的任意破坏性 Token。| 特性维度 | 掩码扩散 (Masked Diffusion) | 均匀状态扩散 (Uniform State Diffusion) | 自回归生成 (Autoregressive) |
|---|---|---|---|
| 噪声表现 | 显式 [MASK] 标记 |
隐式均匀随机 Token 替换 | 无扩散噪声 |
| 去噪动作 | 逐级填充 [MASK] 空间 |
预测全画布，不满足阈值处重填随机噪声 | 逐个 Token 尾部追加 |
| 自纠错能力 | 已固定位置较难回改 | 允许继续修改低置信度位置 | 无法回溯已输出 Token |
| 生成步数 | 取决于采样策略（通常 20-50 步） | 固定或自适应步数（如 48 步） | 与生成文本长度 N 严格一致 |
o^t 和 self-conditioning 向量 S^t，prompt 侧 KV 在整个循环中保持静态。
这张 technical report 图把数据流拆成三部分：prefill prompt、当前 canvas、跨 step 传递的 self-conditioning。工程上最值得看的是这里：Gemma 主干没有被整套推倒重做，改动集中在输入组织、attention mask、KV cache 写入时机和 denoising loop。
diffusion_architecture 图。图里可以直接看到 input query 走 causal encoder，noisy canvas 走 bidirectional denoiser，KV cache 复用，self-conditioning 回到下一步输入。
DiffusionGemma 复用了 Gemma 4 26B A4B checkpoint 的语言主干。公开材料描述的模型是 MoE 架构，30 层，词表约 262K，单 Token 激活约 3.8B 参数和 8 个专家。新增部分主要是 encoder-denoiser 模式切换、self-conditioning 和扩散采样器。
┌────────────────────────┐
│ Prompt: "Write a..." │
└───────────┬────────────┘
│
(Encoder Mode)
Causal Attention
│
▼
┌────────────────┐
│ Static KV │ (Stored in Memory)
│ Cache │
└────────│───────┘
│
(Cross-Attention Guidance)
│
┌───────────────────────────┼───────────────────────────┐
│ Denoiser Loop (48 steps) ▼ │
│ │
│ [Random Canvas] ──► (Denoiser Mode) ──► [Denoised │
│ (256 tokens) Bidirectional Attention Canvas] │
│ │
│ ▲ │
│ └─────────── Self-Conditioning ───────────────┘
└───────────────────────────────────────────────────────┘去噪器通过 self-conditioning 传递上一轮 canvas 状态：
t 步，去噪器输出全画布每个位置的 Logits。t+1 步的输入 Token Embedding 中。DiffusionGemma 仍然有 KV cache，但它的写入时机和普通 AR LLM 不一样：
prompt / 已完成输出 block
-> causal encoder prefill
-> 写入 KV cache
当前 256-token canvas
-> bidirectional denoiser 多步去噪
-> 未完成前不追加 KV cache
canvas 收敛或达到步数上限
-> commit 成 finalized block
-> 再用 causal encoder 路径写入 KV cache
-> 开下一个 256-token canvas“输出的 KV cache”来自 Gemma backbone 在 causal encoder/prefill 模式下对 finalized output block 的重新编码，并非直接复用 denoiser 最后一轮 hidden states。
attention mask 也要分阶段看：
已确定上下文: causal attention
当前 canvas: full matrix / bidirectional attention
canvas 读上下文: 读取 prompt 和已完成 block 的 KV cache
block 之间: autoregressive如果某个 token 48 steps 后仍然不稳定，系统并没有理论保证一定修好。工程上只能接受当前 block、增加 denoising steps、调低温度、改 entropy bound，或者重采样。DiffusionGemma 的速度和质量来自训练好的 denoiser 在常见文本分布上的收敛能力，不是来自“48 步必然正确”的数学保证。
短答案：backbone 尽量保留，生成范式要改。以标准 decoder-only LLM 为例，backbone 通常包括：
embed_tokens、lm_head。DiffusionGemma 保留 Gemma 主干，并增加一条离散扩散推理路径。实现至少涉及这些改动：
C=256 的 noisy canvas。S^t，再并回下一步输入。T 轮“前向 + 采样 + renoise + early stop”循环。| 模块 | 是否保留 LLM backbone | 要改什么 |
|---|---|---|
| Embedding / lm_head | 大体保留 | 继续复用词表与 embedding 空间，但要支持 canvas token 和 self-conditioning 回投。 |
| FFN / MoE / Norm | 基本保留 | 这些是语言能力主干，通常不想动；DiffusionGemma 正是复用了 Gemma 4 的 MoE 主体。 |
| Attention | 结构可复用，行为要改 | 必须支持从 causal 到 bidirectional 的切换，并让 canvas 在每一步读取 prompt 侧上下文。 |
| KV cache | 写入时机变化最大 | AR 里 KV 随 token 逐步追加；这里 prompt 和已 commit block 写 KV，当前 canvas 的 step-to-step 状态主要是 o^t 和 S^t。 |
| Forward API | 必须改 | 输入不再只是 input_ids，而是 prompt context + noisy canvas + diffusion step + optional self-conditioning。 |
| Generate loop | 必须重写 | 从单步 token sampling 改成多步 denoising loop，再加 multi-canvas 外循环。 |
DiffusionGemma 保留 Gemma/MoE 主干，并把生成路径改为去噪：当前 block 读取上下文 KV，并行修正 256 个位置，再用 self-conditioning 跨 step 传递状态。主要改动落在 attention mask、KV 写入时机、训练目标和采样循环，FFN 变化相对较小。
去噪器一次只处理一个 256-token canvas。长文本靠多画布采样：block 内扩散，block 之间自回归。
# 数据流概念伪代码
prompt = "Please write a long sci-fi story..."
kv_cache = encoder.prefill(prompt) # 首次 Prefill 生成初始 KV Cache
while not generate_finished:
# 1. 初始化 256 长度的随机画布
canvas = initialize_random_tokens(length=256)
# 2. 逆向迭代去噪循环 (例如 48 步)
for step in range(num_denoising_steps):
# 融合上一步的概率分布特征
conditioned_embeddings = embed(canvas) + self_conditioning(prev_probs)
# 去噪器进行双向 All-to-All 注意力计算，并 cross-attend 到只读的 kv_cache
logits = denoiser.forward(conditioned_embeddings, kv_cache)
# 采样器决定保留哪些 token，重新对低置信度位置进行 noise 覆盖
canvas, prev_probs = sampler.step(logits, canvas)
if adaptive_stopping.should_stop():
break
# 3. 256 Canvas 去噪完成，识别到结束符或填满
finalized_block = canvas
# 4. 自回归扩展：将 finalized_block 视为 prompt 的延续，送入 Encoder
# 追加并更新静态 KV Cache 缓存。之后该 KV Cache 作为下一画布的上下文。
kv_cache = encoder.incremental_prefill(finalized_block, kv_cache)
if eos_token in finalized_block:
generate_finished = True这就是 block-autoregressive diffusion：一个 block 内用 full attention 反复修，block 完成后再进入 causal KV，作为下一个 block 的上下文。
DiffusionGemma 的采样质量主要受 sampler 和 scheduler 控制。公开示例里常见的配置如下。
为了平衡去噪早期探索和后期收敛，Logits 会被除以一个动态温度 T。
T 采用线性衰减，由 0.8 逐渐降低至 0.4。用于精确控制去噪每一步中保留哪些 Token、丢弃哪些 Token。
"LLM"，熵接近 0）。不一定要跑满 48 步。如果 canvas 提前稳定，可以早停：
为什么 Diffusion 架构适合解决强约束的多变量协同任务（如数独）？
自回归模型 (AR) 顺序预测： DiffusionGemma 并行协同去噪：
┌───┬───┬───┐ ┌───┬───┬───┐
│ 5 │ 3 │ ? ├──► 必须立刻决定第三格 │ 5 │ 3 │ 4 │ 所有格点在 All-to-All
└───┴───┴───┘ 无法考虑未来冲突 ├───┼───┼───┤ 注意力下同时互相制约，┌───┬───┬───┐ │ 6 │ ? │ 8 │ 若后期发现冲突，可通过
│ 6 │ ? │ 8 │ 一旦写错，无法修改 ├───┼───┼───┤ Re-noising 机制在后续去噪
└───┴───┴───┘ │ ? │ 9 │ ? │ 步中将错误擦除并重写。
└───┴───┴───┘vLLM 已集成 DiffusionGemma。下面是 OpenAI 兼容 server 的启动示例：
vllm serve google/diffusiongemma-26B-A4B-it \
--max-model-len 262144 \
--max-num-seqs 4 \
--gpu-memory-utilization 0.85 \
--attention-backend TRITON_ATTN \
--generation-config vllm \
--hf-overrides '{"diffusion_sampler": "entropy_bound", "diffusion_entropy_bound": 0.1}' \
--diffusion-config '{"canvas_length": 256}' \
--enable-chunked-prefill开发者可以使用标准 transformers 库中的 DiffusionGemmaForBlockDiffusion 进行推理：
from transformers import DiffusionGemmaForBlockDiffusion, AutoProcessor
import torch
MODEL_ID = "google/diffusiongemma-26B-A4B-it"
# 1. 初始化处理器与离散扩散专用模型类
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = DiffusionGemmaForBlockDiffusion.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
device_map="auto",
)
# 2. 准备对话模版
messages = [
{"role": "user", "content": "Explain the core difference between AR and Diffusion LLMs."}
]
input_ids = processor.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt"
).to(model.device)
# 3. 运行多画布扩散采样推理
output = model.generate(**input_ids, max_new_tokens=512)
# 4. 打印解码文本
text = processor.decode(output[0], skip_special_tokens=False)
print(text)DiffusionGemma 没有丢掉 KV cache，也不是把整段序列都改成永久 full attention。它把已确定上下文放进 causal KV cache，把当前 256-token canvas 放进 bidirectional denoiser 里反复修。block 内并行带来速度，block 之间的 causal cache 保住长文本上下文。
本项目的 diffusion_engine 是玩具级实现，但有几处概念可以对上 DiffusionGemma：core/attention.py 里的 causal mask 开关对应 encoder/denoiser attention 切换；core/text_conditioning.py 的 prompt cache 对应上下文只读复用；core/scheduler.py 的时间步控制和 DiffusionGemma 的温度调度不等价，但都属于逐步逼近的采样控制。
本章参考了这些资料：