DeepSeek V4 源码解析

• 官方源码：https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/tree/main/inference
• 技术报告：https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/blob/main/DeepSeek_V4.pdf
• 笔者注释：https://github.com/sunkx109/DeepSeek-V4-Flash/tree/master/inference

0. 全局配置速查（V4-Flash）

下表是后文所有性能分析的依据，全部取自 inference/config.json，建议先建立印象：

43 层主结构组成（由 compress_ratios 推出）：

Hybrid 排布：V4-Flash 43 层，前 2 层纯 SWA，之后 CSA/HCA 交错，末尾 1 层 MTP（SWA）；V4-Pro 61 层，前 2 层用 HCA，之后 CSA/HCA 交错。

1. 整体模型结构与推理流程

1.1 模型结构总览

Transformer（model.py）的前向骨架非常简洁：

defforward(self, input_ids, start_pos=0):    h = self.embed(input_ids)                              # [b, s, dim]    h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)       # [b, s, hc_mult, dim]  ← 展开成 4 路并行流for layer in self.layers:                              # 43 个 Block        h = layer(h, start_pos, input_ids)    logits = self.head(h, self.hc_head_fn, ...)            # HC-head 收敛 + LM headreturn logits

一句话概括：embed → 复制成 hc_mult 路 → 堆叠 N 个 Block → HC-head 收敛回 1 路 → 取最后一个 token 算 logits。

几个关键设计点：

• MQA + latent KV：wkv = Linear(dim, head_dim) 把每个 token 压成 head_dim=512 的单个 latent，它同时充当 K 和 V（sparse_attn 里 acc_s = q·kv 用作 K、acc_o = acc_s·kv 用作 V）。所以 KV cache 每 token 只存 1 份 512 维，而不是 K/V 两份。
• Q 低秩：wq_a(dim→1024) → q_norm → wq_b(1024→64×512)，先把 Q 压到 q_lora_rank=1024 再升回多头，降低投影参数量与算力。
• O 分组低秩：输出投影按 o_groups=8 分组做低秩 Linear（见 §2.5）。
• mHC 残差：隐藏状态始终是 4 路（hc_mult=4），每个 Block 在 4 路之间做“混合 → 子层 → 再展开”。详见 §5。

1.2 推理流程（generate.py）

推理是经典的 prefill + decode 两阶段，KV cache 由各层的 self.kv_cache buffer 跨步保留：

1. 预填充（prefill, start_pos = 0）   一次喂入整段 prompt：forward(tokens[:, 0:prompt_len], start_pos=0)   → 压缩 / SWA cache 一次性建好 → 取最后一位 logits → 采样第 1 个新 token2. 逐 token 解码（decode, start_pos > 0）   每步只喂 1 个 token：forward(tokens[:, prev:cur], prev_pos)   # cur = prev + 1   → KV cache 增量更新 → 采样下一个 token → 追加   → 遇到 EOS 或达到 max_new_tokens 停止

几个值得注意的实现细节：

• 采样用 Gumbel-max（sample）：probs / exponential_(1).argmax，等价于多项式采样但避免 torch.multinomial 在 GPU 上的 CPU 同步，更快。
• 左 padding + prompt 覆盖：tokens 用 1 填充，prompt_mask 标记哪些位置是真实 prompt；对于仍在 prompt 内的位置，直接用真值 token 覆盖模型预测（支持 batch 内不等长 prompt 一次性 prefill）。
• 批量推理：generate 支持一个 batch 多条 prompt 同时生成。

prefill vs decode 的本质区别（后文反复用到）：

• prefill：seqlen = prompt_len（>1），压缩器一次处理整段、按 ratio 切块压缩；
• decode：seqlen = 1，压缩器每次只处理 1 个 token，攒够 ratio 个才触发一次压缩。

1.3 一个 Block 的前向（mHC 包裹 attn + MoE）

每个 Block 把 mHC 套在 Attention 和 MoE-FFN 外面：

defforward(self, x, start_pos, input_ids):# ---- 子层 1：Attention ----    residual = x                                        # [b,s,hc,d]    x, post, comb = self.hc_pre(x, hc_attn_fn, ...)     # 4路 → 1路：加权求和    x = self.attn_norm(x)    x = self.attn(x, start_pos)    x = self.hc_post(x, residual, post, comb)           # 1路 → 4路：再展开# ---- 子层 2：FFN(MoE) ----    residual = x    x, post, comb = self.hc_pre(x, hc_ffn_fn, ...)    x = self.ffn_norm(x)    x = self.ffn(x, input_ids)                          # MoE    x = self.hc_post(x, residual, post, comb)return x

也就是说：标准残差 y = x + f(norm(x)) 被替换成了 “hc_pre 收敛 → f → hc_post 展开” 的多路残差。原理见 §5。

2. Attention 详解

V4 的 Attention 分几类：Compress Attention、Sparse Attention、SWA。CSA 同时具备“压缩 + 稀疏”，HCA 只有“极致压缩”，SWA 是保证块内时序因果的旁路。下面分模块讲清每种的计算属性，最后在 §2.5 给出统一的 Attention 计算流程。

2.1 Deepseek Sparse Attention（DSA）回顾

“Sparse” 本质上就是一种减少  计算量的注意力 Pattern——任何通过 mask / 选择跳过部分 KV 计算的模式都算稀疏（SWA 也是）。

DSA 在 MLA 的基础上加了一条旁系分支（Indexer）：先用低维 Q·K 算相关度，从中选出 top-k 的 KV，再让主路 MLA 只对这些 KV 做注意力。

相对传统 MLA，核心差异在 Indexer，其内部有三条计算通道：

• [bsz, heads, seqlen]（每个 KV 的权重）
• [bsz, heads, seqlen, headdim]
• [bsz, seqlen, headdim]

打分流程： → scale → 按头求和 → scale → 得到 [bsz, seqlen_q, seqlen_k] 的相关度，再 top-k。

Indexer 对 QK 的量化（Hadamard + MXFP4，或直接 64-group FP8）见 §4.1。

2.2 Compressed Sparse Attention（CSA, ratio=4）

CSA = 缓存压缩 + 稀疏注意力：先把每  个 token 的 KV 压成 1 条，再结合 DSA，让每个 query 只对  条压缩 KV 做注意力。

核心区别于 DSA：CSA 不只是“选择少算”，而是连 KV cache 的存储都做了压缩（DSA 的 topk 只影响算力、不影响存储）。CSA/HCA 是从 seq_len 维度压缩，而 GQA/MQA/MLA 是从 multi-head 维度压缩——两条正交的优化路线。

压缩怎么做：重叠加权池化

最朴素的压缩是“固定切块、块内加权求和”，但块与块之间会割裂。CSA 用重叠（overlap）窗口：同一个隐藏状态  乘不同权重，得到两组 latent KV（）和对应压缩权重（）；加权时 a 组取当前块 [mi : m(i+1)-1]、b 组取前一块[m(i-1) : mi-1]，使相邻块平滑过渡。

代码里用两个 buffer 暂存“还没攒够一个压缩块”的  和 ，尺寸 [max_batch, 2·ratio, 2·head_dim]：前一个 ratio 段是 overlap 窗口，后一个 ratio 段是当前窗口，每段都存 a、b 两组。

self.register_buffer("kv_state",    torch.zeros(args.max_batch_size, coff*compress_ratio, coff*head_dim, dtype=torch.float32))self.register_buffer("score_state", torch.full((args.max_batch_size, coff*compress_ratio, coff*head_dim), float("-inf"), dtype=torch.float32))

其中 coff = 1 + overlap，overlap = (ratio == 4)。、 的计算：

kv    = self.wkv(x)    # [bsz,seqlen,dim]·[dim, coff*head_dim] = [bsz,seqlen,coff*head_dim]   → C^a,C^bscore = self.wgate(x)  # 同上                                                              → Z^a,Z^b

Prefill 阶段

整段能被 ratio 整除的部分一次性压缩；只有余数部分需要缓存（且余数的后续压缩需要前一块的 ），所以会把“最后一块的 ”暂存进 kv_state 前半段：

if start_pos == 0:                                       # prefill    remainder = seqlen % ratio    cutoff    = seqlen - remainder    ...    kv    = kv.unflatten(1, (-1, ratio))                 # [b, cutoff/ratio, ratio, coff*head_dim]    score = score.unflatten(1, (-1, ratio)) + self.ape   # 加块内位置偏置 APEif overlap:        kv    = self.overlap_transform(kv, 0)        score = self.overlap_transform(score, float("-inf"))    kv = (kv * score.softmax(dim=2)).sum(dim=2)          # 公式12：块内 softmax 加权求和

overlap_transform 把 [b, s, ratio, 2d] 重排成 [b, s, 2·ratio, d]： 放在后半、（当前块的前一块，故有移位）放在前半，使排布刚好满足公式 12：

defoverlap_transform(self, tensor, value=0):    b, s, _, _ = tensor.size()    new = tensor.new_full((b, s, 2*ratio, d), value)    new[:, :, ratio:]  = tensor[:, :, :, d:]   # C^a：当前块    new[:, 1:, :ratio] = tensor[:, :-1, :, :d] # C^b：前一块return new

Decode 阶段

每次处理 1 个 token，只有当 (start_pos+1) % ratio == 0 时才触发压缩；否则只把 、 暂存：

else:                                                    # decode    should_compress = (start_pos + 1) % ratio == 0    score += self.ape[start_pos % ratio]if overlap:        self.kv_state[:bsz, ratio + start_pos % ratio]    = kv.squeeze(1)        self.score_state[:bsz, ratio + start_pos % ratio] = score.squeeze(1)if should_compress:                               # 攒够 4 个 → 压缩            kv_state    = torch.cat([self.kv_state[:bsz, :ratio, :d],    self.kv_state[:bsz, ratio:, d:]], dim=1)            score_state = torch.cat([self.score_state[:bsz, :ratio, :d], self.score_state[:bsz, ratio:, d:]], dim=1)            kv = (kv_state * score_state.softmax(dim=1)).sum(dim=1, keepdim=True)            self.kv_state[:bsz, :ratio]    = self.kv_state[:bsz, ratio:]    # 滚动：当前窗→overlap窗            self.score_state[:bsz, :ratio] = self.score_state[:bsz, ratio:]

压缩完成后，CSA 还会再过一个 Indexer 做 top-k 稀疏选择（见 §2.5）。

2.3 Heavily Compressed Attention（HCA, ratio=128）

HCA 以极致压缩为目标：每 （）个 token 压成 1 条，且不用 overlap、不做 topk 稀疏，纯粗粒度压缩：

kv    = self.wkv(x)    # [bsz,seqlen,head_dim]score = self.wgate(x)if start_pos == 0:                                       # prefill    ...    kv    = kv.unflatten(1, (-1, ratio))                 # [b, cutoff/ratio, ratio, head_dim]    score = score.unflatten(1, (-1, ratio)) + self.ape    kv    = (kv * score.softmax(dim=2)).sum(dim=2)else:                                                    # decode    should_compress = (start_pos + 1) % ratio == 0    self.kv_state[:bsz, start_pos % ratio]    = kv.squeeze(1)       # 直接存，无 overlap    self.score_state[:bsz, start_pos % ratio] = score.squeeze(1)if should_compress:        kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True)

CSA vs HCA 对比：

压缩后的收尾（CSA/HCA 通用）

无论哪种，压缩完都会过 norm、对 rope 维补 RoPE、再量化、写入 KV cache：

kv = self.norm(kv.to(dtype))freqs_cis = self.freqs_cis[:cutoff:ratio] if start_pos == 0else self.freqs_cis[start_pos+1-ratio].unsqueeze(0)apply_rotary_emb(kv[..., -rd:], freqs_cis)               # 压缩后补 RoPEif self.rotate:                                          # Indexer 走这条路：Hadamard + FP4    kv = rotate_activation(kv); fp4_act_quant(kv, fp4_block_size, True)else:                                                    # 主路注意力：非rope维 FP8 模拟    act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True)self.kv_cache[:bsz, start_pos // ratio] = kv.squeeze(1)  # 写入压缩 cache

注意 inplace=True：量化后又反量化回 bf16 写入 cache。所以本参考实现里 KV cache 实际是 bf16 存的，FP8/FP4 量化是为了模拟 QAT 的数值行为；生产引擎会直接保留 FP8 再省一半存储（见 §3.3）。

2.4 SWA（滑窗旁路）

压缩后的 KV 失去了“块内中间 token”的精确因果信息，因此每个 query 额外取最近  个未压缩 KV 参与计算，保证块内时序因果。SWA 的 KV cache 独立于压缩 cache，用循环覆盖写入：

if start_pos == 0:                                       # prefillif seqlen <= win:        self.kv_cache[:bsz, :seqlen] = kvelse:        cutoff = seqlen % win                            # 最后一窗的实际长度# 循环移位：保证最后一窗的 token 排在前面，方便后续按 ring 读取        self.kv_cache[:bsz, cutoff:win], self.kv_cache[:bsz, :cutoff] = kv[:, -win:].split([win-cutoff, cutoff], dim=1)else:                                                    # decode    self.kv_cache[:bsz, start_pos % win] = kv.squeeze(1)

2.5 Attention Block：统一的稀疏计算流程

整体架构：(源自：https://github.com/CalvinXKY/InfraTech)

Attention Block 示意图

无论 CSA 还是 HCA，最终都按 MQA 计算：用压缩后的 KV 同时充当 K 和 V，每个 query 只对 (窗口 KV) ∪ (被选中的压缩 KV) 做注意力。

为统一“CSA 有 indexer / HCA 无 indexer / 都有 SWA 旁路”三种情况，代码对 topk_idxs 和 KV 做了对齐：

# SWA：取当前 Q 最近邻 win 个未压缩 KV 的 indextopk_idxs = get_window_topk_idxs(win, bsz, seqlen, start_pos)if self.compress_ratio:    offset = kv.size(1) if start_pos == 0else win       # 压缩部分在 kv 的后半，offset 是其起始if self.indexer isnotNone:                         # CSA：由 indexer 先压缩后 选 topk        compress_topk_idxs = self.indexer(x, qr, start_pos, offset)else:                                                # HCA：直接取所有压缩 index（不稀疏）        compress_topk_idxs = get_compress_topk_idxs(ratio, bsz, seqlen, start_pos, offset)    topk_idxs = torch.cat([topk_idxs, compress_topk_idxs], dim=-1)topk_idxs = topk_idxs.int()# 真正的稀疏注意力计算：sparse_attn 按 topk_idxs 去 gather KVo = sparse_attn(q, kv_or_cache, self.attn_sink, topk_idxs, self.softmax_scale)

几个尺寸（V4-Flash）需要注意：

• hidden_size=4096，SWA 窗口 128；
• Q 低秩分解维度 1024（这也是 Q 进 indexer 的尺寸）；
• Sparse Attention 计算：Q 64 头，KV 1 头（MQA），head_dim=512；
• Indexer 内部：Q 64 头、head_dim=128，index_topk=512（候选不足 512 按实际、超过只选 512）；
• 主体Attention 有一份 KV cache，Indexer 也有一份独立的 KV cache，两份都走 Compressor 流程

sparse_attn 是一个 FlashAttention 风格的 TileLang kernel：按 topk_idxs gather 出 KV 块、在线 softmax（running max/sum）、并带一个可学习的 attn_sink（attention sink）偏置：

# kernel.py: sparse_attn_kernel 核心循环（节选）for t in T.Pipelined(num_blocks, num_stages=2):# ① 按 topk_idxs 把 KV 块 gather 进 shared mem（无效 idx -1 置 0，注意力分数置 -inf）# ② acc_s = q · kv^T，× scale# ③ 在线 softmax：维护 running max / sum，并 rescale 累积输出 acc_o# ④ acc_o += softmax · kv# 末尾加上 attn_sink，再除以 sum_exp 归一化

Linear Out 分组

输出投影按 o_groups=8 分组做低秩 Linear：把 head 维按每 8 组切 group，组内各自做低秩 wo_a，再 wo_b 合并：

o    = o.view(bsz, seqlen, n_local_groups, -1)wo_a = self.wo_a.weight.view(n_local_groups, o_lora_rank, -1)o    = torch.einsum("bsgd,grd->bsgr", o, wo_a)           # 分组低秩x    = self.wo_b(o.flatten(2))

3. 性能分析：FLOPs 与 KV Cache

这一节回答四个问题：压缩到底省了什么？稀疏到底省了什么？算力降了多少？显存降了多少？

3.1 一句话结论

压缩（Compressor）省的是“KV Cache 存储”；稀疏（topk 选择）省的是“注意力算力 FLOPs”，两者作用在不同维度，互不重叠。

这是理解 V4 注意力效率的钥匙：

• Compressor：把 L 个 token 的 KV 压成 L/r 条 → 存储 ÷ r（CSA ÷4、HCA ÷128）。注意它不减少最终注意力读的“候选条数对应的存储”——它直接减少了物理上要存的条目数。
• Indexer / 稀疏：在压缩后的 L/r 条里只读 top-k 条 → 算力 ÷（约 L/(k+w)）。它不减少存储（被没选中的 KV 仍存在 cache 里，只是本次不算），正如原文所说：“DSA 的 topk select 之后 KV cache 的存储不会发生变化，只会影响 FLOPs”。

3.2 Sparse Attention 的 FLOPs 分析与计算公式

统一公式

设 （Q 头数）、（头维）、=上下文长度、=压缩比、（index_topk）、（窗口）。对单个 query token、单层，注意力计算（ + ）的 FLOPs为：

其中  是该 query 实际读取的 KV 条目数（一次乘加算 2 FLOPs）。代入 V4-Flash 数值：

可见算力与  成正比，算力下降倍率就等于  的下降倍率。各类注意力的 ：

算力下降倍率

定义相对“全量稠密”的下降倍率 。以 decode（每步 1 个 query，最常见）为例，在不同上下文长度下：

读法：128K 上下文下，CSA 层每个 query 的注意力算力只有全量稠密的 1/205，HCA 层只有 1/114。上下文越长，CSA（带 topk 上限 512）的收益越夸张——因为  被 min(512, …) 钉死在 640，而稠密的  线性增长。

相对“只压缩、不稀疏”的额外收益

如果把“压缩”视为已经默认开启（即对比基准是 HCA 那种“压缩但稠密读全部”），那么 indexer 的 topk 在 CSA 层带来的额外算力下降为：

128K 下 CSA：51.2×（即“压缩之上，稀疏再砍 51 倍算力”）。HCA 没有 indexer，它的“稀疏”完全来自压缩本身。

Indexer 的开销是否划算？

Indexer 自己也要算一次（低维）注意力来打分。Attention 其每 query 开销：

128K 下： GFLOPs，此外，还需要加上少许Linear层，总的来说，仅约为同层稠密注意力（17.18 GFLOPs）的 3%。换句话说，花 3% 的算力去选 topk，换回 99.5% 的主路算力节省，极其划算。

3.3 KV Cache 存储分析与 Compressor 的收益

每条目的字节数

主路 KV 每 token = head_dim=512 维 latent（MQA 单份）。本参考实现里量化后反量化回 bf16 存储，所以：

生产引擎若直接保留 FP8（非 rope 维）+ bf16（rope 维 64），则  B，再省约 1.8×。本文按参考实现的 bf16 估算。 Indexer 的 KV 每条 = index_head_dim=128 × 2 B = 256 B。

Compressor 的存储削减（按层）

对单层、长上下文（窗口  相对  可忽略）：

即 CSA 层 ÷4、HCA 层 ÷128。窗口  会略微抬高实际条目数（多存 128 条未压缩），所以 HCA 实际略小于 128。

全模型存储对比（K）

逐层统计 KV cache 条目数：

外加 Indexer 自己的 KV cache（仅 21 个 CSA 层有）：

换算成字节：

注意 HCA 层虽然只有 20 层、却几乎不占存储（每层仅 1152 条），KV cache 的大头是 21 个 CSA 层（占主路 96%）。这也是为什么 CSA 同时配了 indexer——它既贡献最多存储、也最需要稀疏来砍算力。

结论汇总

• Compressor 的存储收益：按层 ÷ 压缩比（CSA ÷4、HCA ÷128）；全模型 128K 下 KV cache 从 ~5.8 GB 降到 ~0.9 GB，约 6–7×（被 indexer 的额外 cache 和 CSA 的大块占用拉低）。
• 若启用 FP8 存储：主路再 ÷1.8，总存储可进一步压到 ~0.5 GB 量级。

4. 其他细节

4.1 量化

结合官方推理代码与 HF 释放的权重，量化出现在以下几处：

量化 kernel 在 kernel.py：act_quant（FP8）、fp4_act_quant（MXFP4，scale 用 IEEE754 位运算取 2 的幂），以及 fp8_gemm / fp4_gemm（带分块 scale 校正的 GEMM）。

4.2 MoE

• 前 3 层是 Hash MoE：维护 [vocab_size, n_activated_experts] 的可学习表 tid2eid，直接按 token_id 索引出要激活的专家——无需打分 topk。
• 其余层是常规 MoE：Gate 打分（sqrtsoftplus）→ bias 偏置选 topk（bias 只影响选谁、不影响路由权重）→ 归一化 → route_scale 缩放。
• Expert 是 SwiGLU FFN（w1 gate / w3 up / w2 down），带 swiglu_limit=10 数值截断；FP4 权重。
• 共享专家：1 个 shared_experts 恒激活，与路由专家结果相加。

路由打分核心：

scores = linear(x.float(), self.weight.float())          # [b*s, n_experts]scores = F.softplus(scores).sqrt()                       # sqrtsoftplusoriginal_scores = scoresif self.bias isnotNone: scores = scores + self.bias    # bias 只影响 topk 选择indices = self.tid2eid[input_ids] if self.hash else scores.topk(self.topk, dim=-1)[1]weights = original_scores.gather(1, indices)             # 用 original_scores 取权重if self.score_func != "softmax": weights /= weights.sum(-1, keepdim=True)weights *= self.route_scale

4.3 MTP

MTP（Multi-Token Prediction）Draft Model 复用主模型的 Block 结构，其 layer_id = n_layers，对应 compress_ratios[43] = 0，即主体是一层纯 SWA 注意力。它在主 hidden state 基础上拼接 token embedding（e_proj + h_proj），跑一遍 Block 后用 hc_head 直接出 logits，用于投机采样草稿。

5. mHC（mixed Hyper-Connection）原理详解

5.1 从标准残差到 Hyper-Connection

标准残差：。信息只有一条流：每个子层的输出叠回到同一条主路上。优点是简单、利于梯度流动；缺点是深层网络里，早期层的特征只能顺着一条道传到后面，容易被不断叠加的新信息“稀释/淹没”。

Hyper-Connection（HC）：同时维护 hc_mult=4 份隐藏状态（4 条并行流）。每个子层前后各加一个可学习的“混合器”：

• 进子层前（hc_pre）：把 4 条流加权求和成 1 条，喂给子层 ；
• 出子层后（hc_post）：把  的输出展开回 4 条流，并与之前的 4 条流混合。

这样信息有了 4 条并行的“高速公路”，早期特征可以搭其中一条直达深层，不必和后续每层的新信号挤一条道。

5.2 hc_pre / hc_post 具体在做什么

记 4 条流为 。

hc_pre（4 → 1）：先对  维做一次线性映射 + RMSNorm，再用 Sinkhorn 解出三组权重 pre/post/comb（见 §5.3）；其中 pre 用来把 4 条流收敛成 1 条：

mixes = F.linear(x.flatten(2), hc_fn) * rsqrt           # [b,s,4d]·[4d,mix_hc] → [b,s,mix_hc]pre, post, comb = hc_split_sinkhorn(mixes, ...)          # pre[b,s,4], post[b,s,4], comb[b,s,4,4]y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=2)  # 4 条流加权求和 → [b,s,d]

hc_post（1 → 4）：用 post 和 comb 把子层输出展开回 4 条流，并混入残差：

# y[c] = post[c]·f(out) + Σ_{c'} comb[c,c']·residual[c']y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2)

直觉上：post 决定“子层新结果分给每条流多少”，comb（4×4 矩阵）决定“旧的 4 条流之间怎么交叉混合”。

5.3 Sinkhorn 归一化：为什么叫 “mixed”

comb 是一个 4×4 的混合矩阵。如果让它自由学习，很容易退化成“某一条流吃掉一切、其余几条闲置”。mHC 的 “mixed” 体现在用 Sinkhorn 算法把 comb 不断行/列归一化，逼近一个双随机矩阵（doubly stochastic）——每行求和为 1、每列求和也为 1：

# kernel.py: hc_split_sinkhorn 核心迭代comb = softmax(comb, dim=-1) + eps          # 先按行 softmaxcomb = comb / (comb.sum(dim=-2) + eps)      # 再按列归一化for _ in range(sinkhorn_iters - 1):         # 反复交替行列归一化（默认 20 次）    comb = comb / (comb.sum(dim=-1) + eps)   #   → 趋近双随机    comb = comb / (comb.sum(dim=-2) + eps)

双随机的物理含义：“公平分配”——每条流贡献出去的总量相等（行和=1），每条流接收进来的总量也相等（列和=1）。这就强制 4 条流都参与、都均衡，防止退化，稳定深层训练。

pre 用 sigmoid（范围 (0,1)）、post 用 2·sigmoid（范围 (0,2)），三者合在一起由同一个线性层 hc_fn 输出，共  个混合系数（mix_hc）。

5.4 为什么 mHC 有用（直观版）

1. 梯度高速公路：4 条并行流给反向传播提供了多得多的回传路径，缓解 43/61 层深网的梯度消失——类似 DenseNet 的密集连接，但混合权重是可学习的。
2. 特征复用与选择性融合：comb 矩阵让每个子层能灵活地把“历史上 4 条流的不同侧面”重新组合，表达力比单条残差强。
3. Sinkhorn 均衡防退化：双随机约束保证 4 条流都不会“饿死”或“独大”，训练更稳、可扩展到更深。

5.5 一个类比

• 标准残差 = 一本笔记本在 43 个人之间传，每人往上加一页；后面的人要翻很厚一叠，早期内容容易被埋。
• mHC = 同时传 4 本笔记本。每个人开工前，图书管理员（hc_pre）把 4 本的摘要抄到一张工作纸上；干完活，管理员再把成果公平地分发回 4 本笔记本（hc_post），并按 comb 在 4 本之间互抄要点。这样早期内容能搭某一本笔记本“直达”后面，而 Sinkhorn 就是要求管理员必须公平分发、不许偏心（行和、列和都为 1）。

末尾的 LM head 也用一个简化版 hc_head（sigmoid 收敛 4→1，无 Sinkhorn）把 4 条流合并成单条再去算 logits。

附：关键文件索引