DeepSeek-V4.1-Flash 结构说明

1总体架构

模型类型 DeepseekV41ForCausalLM,由语言主干(40 层)、ViT 视觉塔(32 层)、 aligner 与 3 层 DSpark MTP 草稿层组成。语言主干被切成两段: L0–L19 是 Causal Encoder,L20–L39 是 Decoder——这是理解全模型所有设计的出发点。

1.2 逐层依赖详图(悬停格子,查看该层的产出与跨层依赖)

每层都自产(不跨层共享):64 头低秩 Q(wq_a→q_norm→wq_b)、SWA 滑窗 KV(128 槽环形,FP8)、 MoE FFN(384 选 6 + 1 共享专家)、HC 混合系数。跨层共享的只有 CSA² 四件套:main KV / index K / Top-512 / 候选池, 且总是由组首的 Full / Reindex 层产出、组内后续层复用——虚线框即共享范围。

2关键参数

权重侧:48 个 safetensors 分片,96,085 个张量,约 475 GiB;dtype 分布 BF16 520 / F32 387 / F8_E4M3 357 / F8_E8M0 47,589 / I8(FP4 打包)47,232。权重名没有 model. 前缀。

3CED:因果编码器—解码器

传统 40 层 Decoder-only 模型,Prefill 时每个 Prompt token 都要完整通过 40 层,每层还要自产 KV。 CED 的做法:Decoder 的 global KV 不由各层自己的隐状态生成,而是统一由 Encoder 末层输出 H₂₀ 经 L20 的 compressor 投影得到:

C = H₂₀ · W^KV  (checkpoint 中只有 L20 一份 compressor 权重,全 Decoder 共享一份缓存)

L21–L39 没有 compressor / indexer 权重,global KV 纯读 L20 的缓存,层间差异由各自的 Q 承担 (存储只需一份,这正是 890 B/token 里 decoder 只占 356 B 的原因)。 于是长度为 S 的 Prompt 只需完整跑前 20 层 Encoder; 只有 Prompt 末尾窗口内的 token 需要继续跑 Decoder 的 SWA Bounded Replay 来构建局部 SWA KV (最后一个 token 走完 40 层,产出第一个回答 token)。

• Prefill:主体计算 ≈ 20/40 = 一半,官方口径每 token 激活 8B 参数;
• Decode:新 token 仍要过 Encoder 产 KV 增量、再过 20 层 Decoder,即完整 40 层(16B/token),CED 不减半 Decode;
• 上下文 4K → 1M(256 倍),单 token Decode FLOPs 只增约 25%(解码计算不再随上下文线性增长;端到端延迟还受 KV 读取、通信、调度影响)。
• SWA Bounded Replay 的另一半用途:decoder 各层的 SWA KV 只能由本层隐状态现算,但窗口只有 128 槽、大小与上下文无关——于是会话挂起/迁移/换 checkpoint 时,SWA KV 直接丢弃不持久化;恢复时只需把最近 128 个 token 重新过一遍(replay),40 层滑窗缓存即重建完毕(过程与近似性见 §4.4)。因此它不占随长度增长的持久化存储(持久化 KV 由此降到上一代 ≈ 1/8)。背后的判据:能不能便宜地重算,决定了要不要存——SWA 能(128 token 重放即可),global KV 不能(重算要把全部历史重跑一遍 Encoder),所以后者才需要 FP4 + 跨层共享这些极限压缩手段。

4Latent Attention 与 SWA 分支

4.1 核心认知:SWA × CSA²,一次注意力 = 两路 KV 拼接

SWA 与 CSA² 不是两种互斥的注意力模式,而是每次注意力同时存在的两路 KV 来源: L2–L39 的每个注意力层,本层自产的滑窗 KV 与组内共享的压缩检索 KV 拼接后,由 Q 统一打分、 同一个 softmax 归一(分母外加 attn_sink 一项)、统一加权求和——一次 sparse_attn 完成:

即每个 query 最多看 128 + 512 = 640 个位置:近处是"本层视角"的细粒度原始 KV, 远处是"组首/encoder 视角"的粗粒度压缩 latent。两路位置有重叠——最近 128 个 token 在两路各有一份 副本(投影不同、不去重,都进同一个 softmax);唯一的可见性限制是组级因果掩码:一个压缩组要等 query 越过组内最后一个 token 才对它可见。

控制这条全局分支粒度的旋钮是逐层配置的 compress_ratios(43 项,含 3 个 MTP)。 它身兼两职:全局分支的开关 + 源层的 token 合并比:

连带影响:压缩缓存条目数 = S/m;压缩条目的 RoPE 取组首 token 位置(第 j 组取 j×m);indexer 打分粒度 = 压缩条目; r=0 的层禁用 YaRN(θ=10000),r>0 的层整层用 θ=160000 + YaRN。注意 m 与 latent attention(通道压缩)、 CSA² 层间共享是三个独立旋钮,别混为一谈。

边界细节(奇数长度怎么办):m>1 时末尾凑不满一组的 token 不硬压——余数落进 compressor 的 kv_state/score_state 状态缓冲,等下一个 token 凑满再池化(组边界是纯位置性的, 一条 latent 可以横跨 prompt/生成边界,如"prompt 末 token + 首个生成 token")。 落单期间它不进压缩缓存、不参与全局分支检索,但必然躺在每层的 SWA 窗内:组最晚在下一个 token 到达时凑满, 而滑窗要 128 步后才把它覆盖——两套缓存有 ~127 步的交接重叠带,任何时刻都不存在"谁都看不见"的 token (新靠 SWA、老靠压缩缓存,交接期两路都在)。

4.2 单层结构:latent attention(layers.N.attn.*)

• Q 低秩:wq_a(5120→1280)→ q_norm → wq_b(1280→32768 = 64×512);
• 单 KV 头:wkv(5120→512)→ kv_norm;每个位置只有一个 512 维 latent,64 头共享,K 与 V 是同一个向量;
• RoPE 只加在末 64 维(qk_rope_head_dim),前 448 维是 NoPE;
• attn_sink:64 维逐头 sink logit(softmax 分母多一项);
• 输出投影分组低秩:输出 reshape 成 8 组(每组 8 头×512=4096),wo_a 块对角(4096→1024)→ wo_b(8192→5120);输出后对 rope 尾维做逆旋转(V 向量带着 RoPE);wo_a 在 checkpoint 是 FP8,官方 convert.py 反量化成 bf16 后按块对角 einsum 使用;
• RoPE 每层只有一组频率:compress_ratio=0 的层(L0/L1/MTP)禁用 YaRN、用基础 rope_theta=10000;compress_ratio>0 的层(L2–L39)整层(Q、SWA 滑窗 KV、压缩 KV 共用)用 compress_rope_theta=160000 + YaRN——名字里的 "compress" 有误导性,压缩 KV 的"位置间隔更大"是靠旋转位置取组首 token(第 j 组取位置 j×m)体现的,不是单独的 theta。

4.3 SWA × 全局分支:两套独立的投影矩阵

SWA 与全局(压缩检索)分支共用 Q 打分与同一个 softmax,但 KV 的投影矩阵完全独立、互不共享—— 同一份隐状态 h,经两套不同的权重产生两种视角的 512 维 latent:

因此"只有 L2/L8/L14/L20 计算并缓存 KV"这个说法只对全局分支成立;SWA 分支每层都在产、都在缓存。 两路缓存在结构上互不读取(对方的 cache 不在自己的数据通路上:wkv 与 compressor.wkv 都只是对 h 的纯投影),但在数值上互相渗透:h 由前序各层堆出, 而每层的注意力输出都融合了两路检索结果——深层的 SWA 条目里早已蒸馏了全局信息,反之亦然。 这个"结构分离、数值交融"的格局,正是下面 replay 敢做近似的前提。

4.4 SWA 的构建(代码级)与 Bounded Replay(部署级)

构建发生在每层的 Attention._window_kv,四步:

1. 投影:kv_norm(wkv(x)) 得 512 维 latent,RoPE 只旋末 64 维;
2. 量化:对 post-RoPE 整向量做 FP8 量化(RoPE 尾维在内);
3. 写入环形缓冲 window_kv_cache[batch, 128, 512](persistent=False,不进 checkpoint): prefill 时序列 ≤128 则整段写入;更长则只取 kv[:, -128:],按 seqlen % 128 旋转对齐填入; decode 时按 start_pos % 128 覆盖单个槽位;
4. 可见性:get_window_topk_idxs 给每个 query 发其因果窗口内的槽位列表, 窗口未满的位置标 -1 屏蔽(prefill 第 k 个 query 只见 k 条;稳态 decode 见整个环,环未填满时旧槽同样标 -1)。

Replay(挂起恢复时的重建):这是部署层的过程——发布的 minimal inference 里没有 replay / restore 函数,只有"多 token 的 chunk 从 start_pos=0 进、单 token 覆盖环槽"两条路径; 真正的重建由 serving 层完成,官方技术报告给出的机制是: 以缓存命中位置 s 为起点、用真实绝对位置重放最近 n_win(=128)个 token, 截断规则为位置 i 的 query 只看 [max(s, i−W+1), i];重放段只重算 SWA KV, 全局压缩 KV 直接复用缓存、既不重算也不覆写(位置基准一旦被重置,持久化的全局缓存与 RoPE 相位就对不上了)。过程要点:

• replay 的 128 个位置输出全部丢弃,唯一目的是把 40 层环形缓冲重新填满;
• 注意力照常两路:SWA 分支处于预热期(第 k 个 token 只见 k 条,其余 -1 屏蔽), 而全局分支读的是持久化的精确压缩 KV——长程信息被锚定,受污染的只有局部纹理;
• 污染随窗口滑动自愈:恢复 decode 后,最老(污染最重)的条目最先被顶出—— 污染程度与剩余寿命精确反相关,128 步内全部换血;
• 代价:128 token × 40 层 ≈ 一次 mini-prefill;恢复会话只需携带 128 个 token id, 不用迁移 2.6 MB 的 SWA 缓存。

5CSA²:压缩检索与跨层复用

5.1 三种静态模式:Full / Reindex / Reuse

每个注意力层仍然计算自己的 Main Query 和 SWA(128)KV,但全局稀疏注意力所用的 main KV、indexer K、Top-K 索引可以跨层共享。三种静态模式逐层指定:

Encoder 侧:18 个 CSA² 层(L2–L19)分 3 组,每组一个 Full 层,压缩率 m=2。 Decoder 侧:仅 L20 为 Full(其 compressor 吃第 20 层注意力子层输入,即 CED 公式的实现), 4 个 Reindex + 15 个 Reuse,m=1。与上一代 CSA 相比,CSA² 删掉了压缩窗口重叠、 压缩阶段的绝对位置编码,indexer K 直接从 main KV 投影,并把 KV 共享与 Top-K 复用解耦。

5.2 分层稀疏索引(Hierarchical Sparse Indexer)

indexer(32 头 × 128 维:indexer.wq_b 1280→4096、wk 512→128、k_norm、weights_proj 5120→32) 若每次扫全部百万 token 自身也会成为瓶颈。因此 Decoder 的第一个 Full 层(L20)扫描全部可见位置时, 除了自己的 Top-512,还按 8 token 一块取最大 indexer score,选出得分最高的 2048 块 = 16,384 个位置 组成共享候选池(candidate_source_layer=20)。后续 Reindex 层只在候选池内打分选 Top-512—— 索引计算量从随上下文线性增长变成近似常数。 实现上,块得分取块内最大 indexer score;含最新位置的未满块会被强制保留(pin 为 +inf), 防止它因位置不满而被老的满块挤出候选池。

5.3 FP4 Main KV Cache 与 compressor

• main KV / 压缩 latent:FP4 E2M1,每 16 通道一个 E4M3 scale(MXFP4),无第二级全局 scale,在 RoPE 之后量化;
• indexer K:FP4 E2M1,每 32 通道一个 E8M0 scale(纯 2 的幂,只用于 Top-K 排序);
• SWA KV 保留 FP8(对整个 post-RoPE 向量量化,RoPE 尾维在内);"长程粗、局部细";indexer Q 推理时同样量化到 FP4;
• 数值依据:归一化后 512 维 latent 的 L2 范数上限 √512,训练期实测最大幅值 ≈ 10,不需要全局 scale;
• SWA Bounded Replay:只 replay 最近 128 个 token,把 SWA 截断在该段内,使 SWA KV 从持久化 cache 中移除(持久化 KV 降到上一代 ≈ 1/8)。

实现上,compressor 的工作方式是:m>1 的层(L2–L19)用 compressor.wkv/wgate(fp32)把 m 个连续 token 经学习到的 softmax 门控池化成一个 latent;m=1 的层(L20–L39)退化为纯投影(bf16,无 wgate)。 compressor 输出的是未加 RoPE 的 latent:indexer 先拿它投影出 index K 写入缓存,然后才对 latent 施加 RoPE、 量化 FP4 写入 main KV 缓存——同一份存储先读后写,两种用途各取所需形态。

6位置编码:部分维 RoPE + 分频段 YaRN

6.1 施加范围:512 维里只有末 64 维旋转

• head_dim 512 = 448 维 NoPE + 64 维 RoPE(qk_rope_head_dim):Q(64 头)与共享 KV latent 都只旋转末 64 维,前 448 维是纯内容维——位置信号与内容信号在通道上分离;
• 旋转方式:相邻两维配成复数,乘位置相位 e^i·p·θᵢ,共 32 个频率 θ_i = base^−2i/64;
• indexer 的 Q/K(128 维)同样只旋尾 64 维,并与主注意力共用同一套频率表(indexer.freqs_cis = self.freqs_cis)——保证检索打分与正式注意力的位置口径一致;
• 输出端有逆旋转:V(= KV latent 本身)带着 RoPE 相位,加权求和后对输出尾 64 维乘共轭相位还原——残差流中的表示因此不含位置相位,缓存也能保持统一的旋转形态(apply_rotary_emb(o, freqs, inverse=True))。

6.2 频率配置:每层一组,共两档

压缩条目的旋转位置取组首 token(第 j 组取位置 j×m,"一个 latent 代表组内第一个 token")—— 压缩 KV 的"位置间隔更大"由此体现,而非单独的 θ;config 里 compress_rope_theta 的名字有误导性, 它实际是"所有带压缩分支的层的 θ"。

6.3 YaRN:按频段分别插值(65,536 → 1,048,576)

对 32 个频率按"该维在训练长度 65,536 内转几圈"分三段(precompute_freqs_cis 的 original_seq_len > 0 分支):

注意:本模型没有采用原版 YaRN 论文的 attention 温度缩放——softmax_scale = 512−0.5 固定不变, 温度补偿可能已在训练中吸收。

6.4 视觉塔:独立的 2D RoPE

ViT 用二维 RoPE:head_dim 64 对半分,前 32 维编码 h 方向、后 32 维编码 w 方向(θ=10,000,fp32 计算),双向注意力; 配对方式是"前后半配对"(chunk 两半)而非语言侧的"相邻两维配对",两套实现各自独立 (vision.py 的 apply_rotary vs model.py 的 apply_rotary_emb)。

7KV Cache 显存账:890 B/token 的推导

每 token 的条目数:Encoder 的 3 组 Full 各按 m=2 摊(3 × 1/2 = 1.5); Decoder 20 层共享 1 份且 m=1(1 × 1 = 1);合计 2.5 条/token:

Encoder 3×½×356 = 534 B  +  Decoder 1×1×356 = 356 B  =  890 B/token

对照 DeepSeek-V4-Flash 的 3514 B/token:890 / 3514 ≈ 25.3%,即 降低约 75%, 同显存理论上可容纳 ≈ 3.95 倍上下文。口径是随序列长度增长的那部分(global KV); SWA KV(每层 128 槽 × 512 维 FP8,40 层 ≈ 2.6 MB/序列)是与长度无关的环形缓冲,不计入 per-token 口径。

8Hyper-Connections(Single-Pass mHC)

残差流不是单向量,而是 hc_mult=4 份并行副本(layers.N.hc_{attn,ffn}_{fn,base,scale},fp32):

• 每个子层前后各一次可学习混合:hc_pre 把 4 份压成子层输入,hc_post 把输出扩回 4 份,并经 comb 矩阵(Sinkhorn 20 次迭代成双随机)混入残差;
• 混合系数由流本身算出:hc_fn(24×20480)+ hc_base(24)+ hc_scale(3);24 = (2 + hc_mult) × hc_mult:pre(4)+ post(4)+ comb(16);投影前对整个展平的 4×5120 流逐 token 算一个 RMS 归一化统计量(不是逐副本);
• Single-Pass:子层算出的系数给下一个子层用(输入混合改用上一层产出的系数),解除"先算系数再混合"的数据依赖:
  X_l+1 = B_l·X_l + C_l·F_l(A_l−1·X_l)
• 部署期融成 Mega-mHC kernel(残差更新 + 输入混合 + 系数预测 + Pre-Norm + FP8 转换),激活内存流量从 (4n+4)d 降到 (2n+2)d(理论下界,n=4),约降一半。

9Engram:参数化 n-gram 记忆

对固定短语、实体名、代码模式等高度重复的局部知识,不必让 Transformer 每次用多层 attention + MoE 重新推导——直接拿最近几个 token 去巨型可训练 embedding 表里查。 MoE 是 Conditional Computation(按隐状态选专家),Engram 是 Conditional Memory (按 n-gram 查记忆);不是 RAG,是端到端训练的内部记忆表。

9.1 计算流程(五步)

1. Tokenizer 压缩:原始 token id 经规范化管道映射为压缩 id(只用于构造查询键,不改模型 tokenizer);管道为 NFKC → NFD → 去口音 → 小写 → 空白合并(含"单空格 token"哨兵保护),129,280 个 token 收敛出 99,092 个唯一 key(config 的 engram_compressed_vocab_size,哈希乘子由它导出,不匹配会静默重哈希整张表);
2. 后缀 n-gram:当前位置取最近 token 构造 2/3/4-gram(不含未来 token,保持因果;n-gram 不跨越图像 span——遇到 dead token 即截断填 pad);
3. 多头哈希:每种阶数 8 个 hash head,共 24 次 lookup,拼接出 24×256 = 6144 维——碰撞在 8 个独立 head 下同时撞上的概率极低;哈希为"奇数乘子 × 压缩 id 后 XOR 滚动"(每层的乘子由 seed=10007×layer_id 的 RNG 生成),每个 (阶数, head) 占有自己质数大小、互不相交的桶区间(从 1600 万起依次取未用过的质数);
4. 投影 + 上下文门控:engram.wkv(6144→25600)投出 4 份 HC 副本各自的 key + 1 份共享 value;门控 = sigmoid(signed_sqrt(归一化点积))——n-gram 提供候选记忆,当前隐状态决定写不写;
5. 门控写入残差流(V4.1 删除了原论文门控后的短因果卷积);图像 token 不参与(engram_mask 关门)。

9.2 本模型配置

收益(官方/原论文):释放早期网络深度与 attention 容量(Multi-Query NIAH 84.2 → 97.0); 参数扩容但不同比增加每 token 计算——一条不同于扩 dense / MoE 的 scaling 路径。

9.3 哈希查表:代码级演示

整个哈希是纯整数运算,输入只有 token id 序列,不碰隐状态——numpy 即可完整复现 (下面对照 inference/engram.py 的 NgramHashState,省略 tokenizer 压缩映射与图像 dead-token 处理;完整可运行脚本见仓库 analysis/ds_v4_1_flash/engram_hash_demo.py)。 这也说明哈希索引完全可以搬去 CPU 预查,GPU 只需要查表结果。

# 配置(config.json): MAX_NGRAM=4, N_HEADS=8, LAYER_IDS=(1,14),
# ENGRAM_VOCAB_SIZE=1600万, COMPRESSED_VOCAB=99092, PAD_ID=2

# ① 桶区间分配: 每层 3 阶(2/3/4-gram) × 8 头,
#    从 1600万起 依次取 24 个互不重复的质数, 各占一段不相交的行区间
primes[layer][order][head]  = find_next_prime(...)   # 如 16000057, 16000079, ...
offsets[layer]              = cumsum(24 个质数)       # 各桶段在表内的起始行号

# ② 乘子: 每层 seed=10007×layer_id 的 RNG 生成, 强制为奇数,
#    上界 int64_max//vocab//2 保证 压缩id×乘子 不溢出 int64
multipliers[layer] = (rng.integers(0, bound, size=4)) * 2 + 1

# ③ 回看窗口: 每个位置取 [当前, 前1个, 前2个, 前3个], 越过序列开头填 pad
tokens = full((L, 4), PAD_ID)
for shift in range(4):
    tokens[shift:, shift] = seq[:L-shift]

# ④ 滚动 XOR: 第 i 步 XOR 完 = (i+1)-gram 的哈希 —— 三种阶数一条链算出
products = tokens * multipliers[layer]        # [L, 4]
rolling  = products[:, 0]
hashes   = []
for i in 1..3:
    rolling = rolling XOR products[:, i]
    hashes.append(rolling % primes[layer][i-1])  # 同一哈希值对 8 个质数分别取模 → 8 路 head

# ⑤ 加桶段 offset → 24 个表内行号, 去 embed 表取 24×256 维
row_ids = concat(hashes) + offsets[layer]     # [L, 24]

实际跑一遍(压缩 id 序列 [5231, 887, 40211, 96, 7712, 330],展开位置 4):

窗口 [当前,前1,前2,前3] = [7712, 96, 40211, 887]
× 层1乘子 [76632096046245, 4839876093313, 35959672319349, 73987337458391]
  ↓ 滚动 XOR
2-gram 哈希 = 7712·m₀                  =  590986724708641440
3-gram 哈希 = 上值 ⊕ 96·m₁             =  590602499734174400
4-gram 哈希 = 上值 ⊕ 40211·m₂ ⊕ 887·m₃ = 2074533117092600990
  ↓ 每个值 % 8 个质数 + offset
24 个行号: [8264676, 16730195, 32156052, ..., 371650938]

• 滚动 XOR 的巧妙处:2/3/4-gram 哈希在一条链上累进,每多 XOR 一项多覆盖一个历史 token,一次扫描 \(O(n)\) 算完三种阶数;
• "8 个 head"不是 8 套哈希函数:每阶只算一个滚动哈希,head 间的独立性来自对 8 个不同质数取模的去相关性(中国剩余定理),向量化一次完成;
• 乘子全为奇数:保证乘法在 int64 环上可逆(不丢比特),且每层用不同 seed → 两层哈希互不相关;
• 图像 token:压缩 id 置为 DEAD=-1,含它的窗口一律填 pad——n-gram 永不跨越图像边界。

10MoE 与路由

• 每层 ffn.gate.weight(384×5120)+ gate.bias + gate.bias_vl(fp32,noaux_tc 校正;图像 span 内 token 用 bias_vl);
• 打分函数 √softplus(x/τ)(sqrtsoftplus,τ=gate_temp=1);bias 只参与选专家、不进权重;权重 = 原始分数归一化 × 1.5(routed_scaling_factor);
• 路由专家 SwiGLU(w1/w3:5120→2304,w2:2304→5120),FP4 E2M1 量化;1 个共享专家(FP8)每 token 必过;
• 专家前向带 swiglu_limit=10 的 clamp(来自训练侧 fp8/fp4 激活值域约束);up 分支(w3)双边 clamp 到 ±10,gate 分支(w1)只封上界(≤10),与训练 kernel 一致;
• MTP 层是缩版 MoE:128 个路由专家选 3。

11DSpark MTP(推测解码草稿)

num_nextn_predict_layers=3,前缀 mtp.0/1/2(layer_id 40/41/42),每层是完整 Block (DSparkAttention + 缩版 MoE + HC),embed/head 与主干共享。目标:每个 decode 步,一次前向产出 block_size=5 个草稿 token + 各自置信度,交主干并行验证。注意开源 repo 只实现前向 (forward_spec),accept/reject 的验证循环不在发布范围内(官方注释: "the speculative-decoding loop itself is out of scope for this repo")。

11.1 一步 decode 的数据流总览

11.2 三层内部结构:哪些权重各有几份

即三层 Block 主体同构,差异只在首尾:mtp.0 多一条"主干 → 草稿流"的投影通路, mtp.2 多输出头三件套。中间层 mtp.1 是纯计算层,只加深草稿流的深度。

11.3 逐步走一遍:forward_spec 的执行顺序

1. 收集:主干 forward 途中,target 层(37/38/39)块入口各存一份 h.mean(dim=2),末尾 cat 成 [15360];
2. 入口(mtp.0.forward_embed):main_x = main_norm(main_proj(main_hidden)); 同时构造草稿输入 [t, noise, noise, noise, noise](noise_token_id=128799),过共享 embed 后复制成 4 份 HC 流;
3. 三层 Block:每层 DSparkAttention 先用 main_x 经本层 wkv 现算 main_kv(RoPE 取主干位置 t,FP8 量化),写入本层自己的 128 槽滑窗环 start_pos % 128; 草稿 Q/KV 用未来位置 t+1…t+5 的 RoPE;然后 MoE、HC 与主干 Block 完全相同(草稿纯文本,不传 image_mask、不用 bias_vl);
4. 输出头(mtp.2.forward_head):hc_pre 合流 → norm → 共享 head 得 5 位 logits; 逐位循环:markov_head(已采 token)产出 logit 偏置与 256 维 markov_embed,偏置加到当前位 logits 后采样出下一位—— 即草稿块内部依赖由马尔可夫链在头部串行化,而不是靠注意力掩码;
5. 置信度:合并后 hidden 与 5 个 markov_embed 拼接(5120+256)经 confidence_head.proj(fp32)得每个草稿 token 的置信度,供验证阶段决定接受前缀长度;
6. prefill 特例:start_pos=0 时三层只对 main_x 做 wkv seed 滑窗环,forward_spec 直接返回 None,不产生草稿。

11.4 DSparkAttention 的可见性与缓存细节

• 可见范围(get_dspark_topk_idxs):整个 128 滑窗环(最旧在前)⊕ 本块全部 5 个位置—— 块内 5 个位置互相可见、无因果掩码,因为草稿是并行生成的,先后依赖交给头部的马尔可夫采样处理;
• 两组位置口径:滑窗环里存的是主干信息流的 KV(main_x 的投影,RoPE 位置 = 主干真实位置); 草稿自身的 KV 只占未来位置 t+1…t+5,当步用完即弃,不进任何缓存—— 所以 MTP 的持久状态也只有 3 个 128 槽环形缓冲,同样与上下文长度无关;
• 纯 SWA 层待遇:compress_ratio=0,θ=10000 无 YaRN;RoPE/逆旋转、分组低秩输出(wo_a 块对角 → wo_b)与主干注意力逐行同构;
• markov_head:embed/head 各 129280×256(rank 256),本质是一张低秩分解的一阶马尔可夫转移表, 以加性 logit 偏置注入——机制详解见 §11.5。

11.5 markov_head:三个关键点

① 本质是查一张一阶马尔可夫转移表,实现上做了低秩压缩。 理想形态是 V×V 的 bigram logit 表 B(B[x,y] = "上一位是 x、下一位是 y 的偏置"),V=129,280 时完整表 ≈ 33 GB; 实现压成 rank-256 分解 B = E·Hᵀ(两张 129280×256,bf16 ≈ 132 MB,约 250 倍压缩)。 查询时只需:上一位 token id → 取 E[x](256 维)→ 乘 H → 得到整行 B[x,:]。 只查 id、不读隐状态——这正是"一阶"的含义,也是它几乎零成本的原因。

② 串行链条:注意力并行算 5 位 logits,先后依赖由头部 5 步循环补上。 每一位的最终分布由两路相加得到:

最终 logits = head(隐状态) + markov_head(上一位 token id)
↑ 深层上下文理解     ↑ 浅层词序习惯

# DSparkBlock.forward_head(mtp.2);logits [B,5,V] 由共享 head 一次算出
output_ids[:, 0] = input_ids                               # 首位 = 主干刚采的真实 token
for i in range(block_size):                                # 逐位串行 0…4
    bias, mk_emb = self.markov_head(output_ids[:, i])      # 拿上一位 id 查转移表
    logits[:, i].add_(bias)                                # 偏置加到第 i 位 logits
    output_ids[:, i+1] = sample(logits[:, i], temperature) # 采出第 i+1 个草稿,成为下一步的键

这就是草稿块内可以不用因果掩码的原因:并行归注意力,串行归这个循环。 (副产品 mk_emb 逐位攒齐后拼进 confidence_head 输入。)

③ 为什么能直接相加:logit 空间相加 = 概率空间相乘。 softmax(z+b) ∝ softmax(z) ⊙ eᵇ,即"主干的概率分布 × 马尔可夫先验分布,再归一化"—— 贝叶斯式的"先验 × 似然 → 后验",也是经典的 Product of Experts: head 提供深层上下文证据,markov_head 提供局部词频证据,两个独立信息源的对数概率本来就该相加。 且偏置幅度是端到端训练校准的;再加上它只作用于草稿、主干验证兜底, 最坏情况只是接受率下降,不会影响最终输出的正确性。

注:checkpoint 配置 block_size=5、target 层 {37, 38, 39}(对应 main_proj 输入 15360 = 3×5120); 推理代码 __main__ 冒烟示例用的是 block_size=6、target (3,4),仅为演示参数,勿当作模型配置。

12视觉塔(ViT + aligner)

12.1 数据流总览:像素 → patch → ViT → aligner → 文本流

记号对应:vision.* 为 ViT(含 patch_embed 与 32 层 Block),aligner.* 为跨模态投影, image_start/end/newline 挂在 embedding 侧。位置信息的处理方式见 12.4。

12.2 图片尺寸规划:token 预算、分辨率上限与超限处理

约束对象要分清:vision_max_n_token=1024 限制的是经 aligner 3×3 折叠后进入 LLM 的 token 数, 不是 ViT 侧的 patch 数。一格 LLM token 对应 42×42 像素(patch 14 × downsample_ratio 3), 即 n_llm_h × n_llm_w = ceil(H/42) × ceil(W/42),预算口径(含换行与边界):

num_image_tokens = n_llm_h × (n_llm_w + 1) + 2 ≤ 1024

式中 +1 是每行末尾各一个 image_newline(逐行插入,行数越多开销越大),+2 是 image_start/end。所以即使折叠后只有 1×1 格,序列也是 [start][IMAGE][newline][end] 共 4 个 token,而不是"1 + 前后 2 个"。