DeepSeek-V4-Flash 架构深度拆解

版本：v0.1 草稿 · 撰写日期：2026-06-08 · 范围：V4-Flash（284B/13B 激活）

CH 0. 摘要与阅读路径

V4-Flash 是 DeepSeek 2026-04-24 发布的开源 MoE 模型。本报告按 9 章展开：

• CH1：V3.2 → V4 演进脉络
• CH2：V4-Flash 整体架构
• CH3-6：CSA+HCA 注意力、MoE、mHC、Muon 四大重点
• CH7：1M 上下文 / FP4+FP8 / 后训练三支撑项
• CH8：源码映射汇总
• CH9：总结

适合：先看 CH 0–2 拿全貌，再按兴趣挑 CH3–6；CH8 适合边读边查。

CH 1. V3.2 → V4 演进脉络

1.1 V3.2 关键架构回顾

V3.2 是 V4 之前 DeepSeek 的旗舰基座（2024-12 发布），由 671B 总参 / 37B 激活参数构成，对外提供 Base 与 Instruct 两套权重。其架构由四大支柱组成：

• MLA（Multi-head Latent Attention）：把 K/V 各自先压到 d_c=512 的潜空间再投影到多头，推理时只需缓存 d_c 维潜在向量即可恢复多头 K/V。这是 V3 系列相对 MQA/GQA 的关键省显存设计。
• DeepSeekMoE：在 V2 的细粒度专家基础上叠加共享专家隔离——每个 MoE 层包含 1 个 shared expert（始终激活）与若干 routed experts（top-k 选择），兼顾通用知识与专项能力。
• FP8 混合精度训练：在 H800 上以 E4M3 / E5M2 为主的 FP8 路径完成 dense forward / backward，是 V3 系列能在受限算力下扩展到 671B 的工程关键。
• Aux-Loss-Free 负载均衡（V3.2 引入）：放弃传统 aux loss 拉均衡，改用 per-expert 标量偏置 b_i，命中率高 → b 减小，命中率低 → b 增大；b 不进梯度、不影响 routing weight，仅影响 top-k 选择。

在 128K 上下文与中短推理任务上，这套组合的"质量 / 成本"比几乎是当时开源 SOTA。但 V4 的设计文档明确指出，V3.2 在 1M 长上下文与超大规模训练上还有四类具体瓶颈（详见 §1.2）。

1.2 V3.2 的四个瓶颈

V3.2 的四个瓶颈与 V4 的四处创新形成严格一一对应：

1. 长上下文 FLOPs 灾难：MLA 把 KV cache 压到 d_c=512，但 Attention 计算仍是 O(T²)。T=1M 时单 token 推理 FLOPs 仍随序列长度二次增长，远超在线服务预算。
2. KV cache 仍线性增长：即使 d_c 小，cache 总量 = T × d_c × L × 2(QK) = 1,048,576 × 512 × 61 × 2 ≈ 6.55 × 10¹⁰ float，BF16 下单样本约 131 GB（不是 240 MB）——100 并发即 13.1 TB 显存，1M 上下文工程化困难。
3. 训练稳定性问题：V3.2 训练到中后期，深层 block 的 ||x_{l+1} − x_l|| 收敛到一个非零常数附近，导致梯度反向传播时信号被持续放大/缩小，需要仔细调节 lr 和 init_std 来避免损失曲线震荡。V3 论文 §5.1 实际采用 Pre-Norm（与 V4 相同）+ DualPipe 流水线 + 通信-计算 overlap 等工程手段缓解稳定性问题；V3.2-Exp 进一步引入 aux-loss-free bias（per-expert 标量偏置 b_i）让路由负载均衡更稳定。（V3 paper 公开训练量为 14.8T tokens；本节瓶颈描述基于 V3.2-Exp 公开报告的定性结论，具体训练量未官方公开）
4. 优化器收敛速度：AdamW 在 MoE 路由稀疏梯度上效率受限；V3.2 训练耗时以月为单位（具体数字未官方公开），单次实验调参周期长、成本高。V4-Flash 公开训练量约 32T tokens，是 V3 公开 14.8T 的 2.2×，没有更快、更稳定的优化器无法在工程预算内完成。

这四类瓶颈直接催生 V4 的四处创新：(1)(2) → CSA + HCA 混合注意力；(3) → mHC 多通道残差；(4) → Muon 正交化优化器。

1.3 V4 产品线划分

V4 系列按"参数规模 × 训练 / 推理形态"两个维度切成 4 个模型卡：

切分逻辑有三层：首先，Base 与 Instruct 的差别是 post-training 流水线（V4 走"领域专家 → on-policy 蒸馏"，见 §5）；其次，FP8 与 FP4+FP8 的差别在于 routed expert 权重是否量化到 FP4——FP4 路径需要 QAT，不能在 Base 权重上直接启用，故只有 Instruct 提供 FP4 变体；最后，Pro 与 Flash 是参数规模切片——Pro 走"49B 激活 × 1.6T 总参"的旗舰路径，Flash 走"13B 激活 × 284B 总参"的高性价比路径，单 batch 在 2×H200 / 8×H100 上即可跑。

1.4 V4-Flash 定位

V4-Flash 是本报告全程追踪的目标模型，有五点关键定位：

• 13B 激活：在 2×H200 / 8×H100 上单 batch 可跑；本报告全部源码引用（inference/model.py / inference/kernel.py）与图示都以 V4-Flash 为实例——其实际超参为 43 层、hidden_size=4096、64 个 attention head、1 个 KV head（MQA）。
• FP4 专家：256 个 routed expert 权重以 FP4（float4_e2m1fn_x2，block=32，E8M0 scale）存储，attn/shared expert/公共部分仍为 FP8。FP4 路径在当前硬件上 peak FLOPs 与 FP8 持平，但显存占用与未来硬件（如 FP4 原生 tensor core）上有显著优势。
• 1M 上下文：默认 max_position_embeddings=1,048,576，RoPE 用 YaRN 从 64K 扩展（factor=16，β_fast=32，β_slow=1）；每层另加 sliding_window=128 的局部窗口。
• License：MIT，权重 + 推理代码全部开源（仓库 DeepSeek-AI/DeepSeek-V4-Flash）。
• 三个推理模式：Non-think（直接出答案）/ Think High（中等推理预算）/ Think Max（最大推理预算，推荐 384K context window）。V4-Flash-Max 是 Think Max 模式下的同源衍生卡。

围绕 V4-Flash，CH2 将展开"整体架构"——其余章节（CH3–CH6）分别深入 CSA+HCA 注意力、MoE、mHC、Muon 四大重点。

CH 2. V4-Flash 整体架构

CH1 已交代 V4 的"为什么"——四大瓶颈 → 四大创新。本章回答"是什么"：V4-Flash 究竟是哪种形态的模型？CH3–CH6 将分别深入四大重点模块。

2.1 V4-Flash 超参数表

V4-Flash 的全部关键超参来自仓库 config.json（HF 风格）与 inference/config.json（仓库内部风格）两份文件。下表整合两份配置，标出每个字段在源文件中的真实取值：

解读：与 V3.2 相比，V4-Flash 的核心收缩有四：(1) 激活参数 37B → 13B（×0.35），单 batch 在 2×H200 / 8×H100 上即可跑；(2) Block 数 61 → 43（×0.70）；(3) 隐藏维度 7168 → 4096（×0.57）；(4) top-k 8 → 6（×0.75），shared expert 仍为 1。基线对比：

V3.2-Exp 基线（671B/37B 激活）：61 层 × hidden 7168 × 64 head × 256 expert × k=8
V4-Flash 目标（284B/13B 激活）：43 层 × hidden 4096 × 64 head × 256 expert × k=6

V4-Flash 在大幅压缩"每层 FLOPs"的同时新增三个 V3.2 没有的设计维度：每层可切换的 CSA/HCA 注意力压缩（compress_ratios 43 项）、4 通道 mHC 残差（hc_mult=4）、以及 FP4 量化专家。这一组参数使"43 层 × 13B 激活 × 256 专家 × 1M 上下文 × FP4 专家"成为合理组合——压缩比/通道数/量化粒度三档互相配合，把单 token 推理 FLOPs 与 KV cache 都压到 V3.2 的约 1/3–1/4。下一节给出 V4-Flash 单个 Block 的数据流框图。

2.2 V4-Flash 顶层框图（单个 Block）

V4-Flash 单个 Block 的数据流可概括为"双分支注意力 → 残差混合 → MoE → 残差混合"四步。设输入张量 x ∈ ℝ^{T×d}，d=4096：

x ∈ ℝ^{T×d}
  ↓
RMSNorm (Pre-Norm)
  ↓
┌─ CSA branch (CH3) ─────────────────────┐
│  x → Compressor (m=4) → Q,K,V         │  → o_csa ∈ ℝ^{T×d}
│       ↓                                │
│  Indexer (64 head × 128 dim, FP4)       │
│  → top-k=512 (CSA 专属评分)            │
│  → Sparse Attention (FlashAttn-style)  │
└────────────────────────────────────────┘
┌─ HCA branch (CH3) ─────────────────────┐
│  x → Compressor (m=128) → Q,K          │  → o_hca ∈ ℝ^{T×d}
│       ↓                                │
│  Position top-k=512 (无 Indexer)        │
│  → Sparse Attention                    │
└────────────────────────────────────────┘
  ↓
o = α · o_csa + (1-α) · o_hca
  ↓
mHC 残差混合 (CH5, hc_mult=4, Sinkhorn-Knopp)
  ↓
RMSNorm (Pre-Norm)
  ↓
┌─ MoE (CH4) ───────────────────────────┐
│  Gate (sqrtsoftplus + aux-loss-free b)  │
│  → top-6 routed + 1 shared             │
│  → 256 routed experts (FP4, block=32)   │
│  → 1 shared expert (FP8)                │
│  → 加权和 × 1.5                         │
└────────────────────────────────────────┘
  ↓
mHC 残差混合
  ↓
→ 下一个 Block (× N=43)

需要注意三处与简化图示的差异：首先，CSA 与 HCA 在源码中共享同一组 MQA-KV（num_key_value_heads=1）和同一 Compressor 算子，仅 compress_ratios[layer_id] 决定本层走 CSA (m=4) 还是 HCA (m=128)；其次，实际 V4-Flash 走"逐层交替"模式：前 2 层（compress_ratios[0]=0, [1]=0）是纯滑窗（compress_ratio=0 表示该层完全跳过 Compressor，走纯 window attention），末尾 compress_ratios[42]=4（CSA），compress_ratios[43]=0 是 list 占位符、不对应任何层；第三，每个 Block 在 attention 输出与 MoE 输出后各做一次 mHC 残差混合（hc_pre 与 hc_post 两次），所以 43 层共有 86 次 mHC Sinkhorn-Knopp 投影，每次 20 迭代。下一节用一张图把"一个 token 从字符串到下一个 token"的全生命周期画出来。

2.3 Token 生命周期

一个 prompt 从字符串到生成下一个 token，经过 6 个阶段。下表给出每步的关键代码入口（路径以 _work/hf-snapshot/ 为根）：

阶段 1 的 thinking_mode 取值为 "chat"（非思考模式）或 "thinking"（开启思考），后者还可叠加 reasoning_effort="high" 或 "max" 切换强度（详见 §2.4）。阶段 4 的 Transformer.forward(input_ids, start_pos) 内部按 compress_ratios[layer_id] 决定每层走 CSA、HCA 还是纯滑窗；KV cache 实际分配在 Attention.__init__ 中按 kv_cache_size = window_size + max_seq_len // ratio 计算（model.py:L473）。阶段 6 的 sample 函数用 Gumbel-max trick 替代 torch.multinomial 以避免 GPU-CPU 同步，配合 temperature=1.0, top_p=1.0 即 V4 官方 README 推荐的默认采样参数（Think Max 模式推荐 context window ≥ 384K tokens）。

整个 6 步流程串成下图的时序。CH3–CH7 将逐一展开"Block 内部"4 个子模块（CSA / HCA / MoE / mHC）以及 Muon / 量化 / 上下文扩展 3 个支撑项的具体实现。

2.4 三种推理模式

V4-Flash 与 V4-Pro 一样支持三种推理 effort 模式（README 表格 + encoding_dsv4.py 双重定义）：

切换方式在源码里通过两个独立参数控制：

1
2
3
4
5
6
7

# encoding/encoding_dsv4.py L506
def encode_messages(
    messages: List[Dict[str, Any]],
    thinking_mode: str,            # "chat" 或 "thinking"
    reasoning_effort: Optional[str] = None,  # "max" / "high" / None
    ...
) -> str: ...

三个模式的本质差异有两层：第一层，thinking_mode="thinking" 时，prompt 中插入 <think> 起始 token 并引导模型先输出完整 reasoning；"chat" 时直接给答案。第二层，当 reasoning_effort="max" 时，render_message（L262）会在 system prompt 前额外插入 REASONING_EFFORT_MAX 前缀（来自 encoding_dsv4.py L86 的指令文本：“If thinking_mode is enabled … you MUST output your complete reasoning …"），告诉模型"必须把全部 reasoning 写在 内、再做 final answer”。OpenAI 兼容 API 调用示例（伪代码）：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from openai import OpenAI
client = OpenAI(base_url="https://api.deepseek.com", api_key="...")

# Non-think
resp = client.chat.completions.create(
    model="deepseek-v4-flash",
    messages=[{"role": "user", "content": "1+1=?"}],
    extra_body={"thinking_mode": "chat"},
)

# Think Max (需要 context window ≥ 384K)
resp = client.chat.completions.create(
    model="deepseek-v4-flash",
    messages=[{"role": "user", "content": "证明 x^2 + y^2 ≥ 2xy"}],
    extra_body={"thinking_mode": "thinking", "reasoning_effort": "max"},
    max_tokens=32768,
)

实际仓库 inference/generate.py 第 125 行用的就是上述编码流程：tokenizer.encode(encode_messages(messages, thinking_mode="chat"))。CH3 接下来展开"双分支注意力"——CSA 与 HCA 的具体设计。

CH 3. 注意力机制：CSA + HCA

CH1 已说明 V3.2 在 1M 上下文下的两个根本瓶颈：(1) Attention 计算量随 T² 增长；(2) KV cache 随 T 线性增长但仍占大量显存。V4 用 CSA + HCA 混合注意力在逐层交替的形态下同时解决这两个问题。本章按"为什么 → 怎么算 → 怎么配合 → 实际省多少 → 源码在哪"的顺序展开。

3.1 标准 Attention 在 1M 上下文下的复杂度灾难

标准 Multi-Head Attention 的计算形式化如下。设输入序列长度 T、隐藏维度 d、头数 h、单头维度 d_h = d / h：

Q, K, V ∈ ℝ^{T × d}                # 三个投影后的张量
Q, K, V split into h heads: ℝ^{T × h × d_h}
Attention(Q, K, V) = softmax(QK^⊤ / √d_h) · V

FLOPs 分析（单层、单 token decoding 视角）：

• QK^⊤：每个 query 算 T 个内积，单次内积 d_h 次乘加 → T · d_h FLOPs
• softmax(QK^⊤)：标量函数，复杂度 O(T)，可忽略
• softmax · V：每个 query 算 d_h 个加权求和 → T · d_h FLOPs
• 合计：2 · T · d_h = 2 · T · d（多头求和后头维被 h × d_h = d 吃回）

对 V4-Flash：d = 4096。T = 1,048,576（1M） 时，单层单 token 推理 FLOPs ≈ 2 × 1,048,576 × 4096 ≈ 8.6 × 10⁹。乘以 43 层 + 注意力之外的 QKV/O 投影 + MoE，单 token 1M 上下文 FLOPs 仍以 T² 量级主导。

为直观看到 T² 的灾难性，列出预填充（prefill） 视角下的整张 attention 矩阵（一次算完整 prompt）：

• 矩阵尺寸：T × T = 1,048,576 × 1,048,576 ≈ 1.1 × 10¹² entries
• 在 causal mask 下，仅下三角（含对角线）非零；可参与 attention 的 entries 数量 = T(T+1)/2 ≈ 5.5 × 10¹¹
• 单层 attention FLOPs ≈ 5.5 × 10¹¹ × d_h × 2 ≈ 4.5 × 10¹⁵
• 43 层累计 ≈ 1.9 × 10¹⁷ FLOPs
• 以 H200 (~1979 TFLOPS FP16 dense) 跑预填充，需 1.9×10¹⁷ / 1.98×10¹⁵ ≈ 96 秒（单卡单 batch 上界，忽略 IO）

KV cache 分析：

• 每层存 T × d × 2 个 float（K 与 V 各一份）
• V4-Flash：1,048,576 × 4096 × 2 ≈ 8.6 × 10⁹ float/layer
• 43 层累计 ≈ 3.7 × 10¹¹ float ≈ 1.4 TB（FP16/BF16 单精度）
• 即便换 V3.2 的 MLA（把 K, V 压成 d_c=512 的 latent），仍需 T × d_c × 43 × 2 ≈ 4.4 × 10¹⁰ float ≈ 168 GB

结论：纯标准 attention 在 1M 上下文下既算不动（T² FLOPs 主导），又存不下（cache 百 GB 量级）。MLA 把 cache 压下来但没压计算量，CSA + HCA 的核心创新是把"计算稀疏化"与"cache 稀疏化"同时做掉——下一节展开 CSA。

3.2 CSA（Compressed Sparse Attention）原理

核心思想：先把 K, V 沿时间维做有重叠的 softmax-pool 压缩（m=4，每 4 个 token 合成 1 个压缩向量），再用 Indexer 评分为每个 query 选出 top-k 个压缩位置做稀疏 attention。

形式化（公式 3.1-3.3）：

(3.1) 压缩 K：设共享 Compressor 为 CComp，对输入 x ∈ ℝ^{T × d}，输出

c_K = CComp_K(x) ∈ ℝ^{T/m × d}        # m=4 (CSA)
c_V = CComp_V(x) ∈ ℝ^{T/m × d}        # m=4 (CSA)

关键点：CSA 的 m=4 时，Compressor 走重叠窗口（overlap = True）——同时维护 2m=8 个 token 的滑动窗口，把"前 4 token + 后 4 token"两段拼起来做 softmax-pool，边界处的压缩更平滑。

(3.2) 索引分数：设 query 投影 q_t ∈ ℝ^{d}、Indexer 专属 key k_idx ∈ ℝ^{T/m × d}、权重 w_t ∈ ℝ（每 query 一个标量）：

I_t = relu(q_t · k_idx^⊤) ∈ ℝ^{T/m}          # query-key 内积
score_t = I_t * w_t ∈ ℝ^{T/m}                 # 逐头乘权重
top-k_idxs_t = topk(score_t, k)               # k=512

Indexer 内部细节（CSA 专属）：

• 独立 Compressor（Hadamard 旋转 + FP4 模拟）——与 Attention 共享的 Compressor 是两个不同实例
• Hadamard 旋转（rotate_activation）：把 d 维向量与一个固定的 Hadamard 矩阵相乘，等价于无参数的随机正交变换，作用是打散信息让量化更鲁棒
• FP4 模拟（fp4_act_quant）：用 float4_e2m1fn_x2 把 q 和 k 量化到 4 bit，但仍跑 FP16 计算——这是 QAT（Quantization-Aware Training） 的标准做法，让前向分布贴近部署形态

(3.3) 稀疏 Attention：

o_t = softmax(q_t · c_K^⊤ / √d) · c_V        # 仅在选中的 k 个位置上算

复杂度分析：

• 标准：O(T² · d)
• CSA：O(T²/m + T·k) = O(T²/4 + T × 512)，其中 T²/4 来自"压缩后 query 数 × 压缩后 key 数"，T·k 来自 top-k 选择后实际算的 attention 项数
• 当 T = 1M：(10¹²/4) + (1.5×10⁹) ≈ 2.5×10¹¹ 主导项是 T²/4，仍比标准 attention 小 4 倍

Cache 节省：

• CSA 存的是 c_K, c_V 而非原始 K, V
• shape 从 [T, d] 变成 [T/4, d]——cache 减到 1/4
• 但完整 cache仍包含 128-token 滑窗（window_size=128，与 V3.2 MLA 同源），所以实际每层 cache = window_size + T/4 = 128 + 1,048,576/4 ≈ 262,272 个向量

关键实现（代码 3.6 节给出完整片段）：

• Compressor.forward（inference/model.py:L316-L378）：用 wkv、wgate 两个 Linear（fp32）做 K, V 的 gating，softmax 权重来自 ape（learned absolute position embedding）+ wgate(x)
• Indexer.forward（inference/model.py:L402-L434）：先 wq_b 投影 q，再 rotate_activation + fp4_act_quant，最后 index_score = einsum(q, kv_cache) → relu → × weights → topk

下一节看 HCA——CSA 的"更激进压缩"版本。

3.3 HCA（Heavily Compressed Attention）原理

核心思想：HCA 是 CSA 的"超压缩"变体——把压缩比从 m=4 拉到 m'=128，同时砍掉独立 Indexer，改用位置 top-k。一句话：与其"压缩 KV → 评分 → 选 top-k"，不如"重度压缩 KV → 直接按位置选"。

为什么这样设计？ 压缩比 128 已经把序列长度压到 T/128 = 1,048,576 / 128 = 8,192，此时"全部 8,192 个压缩 key"对单 query 算 attention 的 cost ≈ 8,192 × d = 3.4 × 10⁷ FLOPs，本来就不大，再做 Indexer 评分反而引入额外开销。V4 的取舍是：HCA 让"压缩已经足够稀疏"成为前提，省掉 Indexer。

形式化（公式 3.4-3.6）：

(3.4) 重度压缩 K, V：

ĉ_K = CComp_K(x) ∈ ℝ^{T/m' × d}    # m'=128 (HCA)
ĉ_V = CComp_V(x) ∈ ℝ^{T/m' × d}    # m'=128 (HCA)

关键差异（vs CSA）：

• m'=128 比 CSA 的 m=4 大 32 倍
• overlap = False（HCA 不开重叠窗口，公式上 coff = 1 + 0 = 1）
• 用同一 Compressor 类但不同 ratio 参数——源码 Compressor(args, compress_ratio=128, ...)（inference/model.py:L467）

(3.5) 位置 top-k（无 Indexer）：

topk_idxs_t = positions[t · (T/128 / T) : t · (T/128 / T) + 512]    # 等间距切

实际源码是 get_compress_topk_idxs(ratio, bsz, seqlen, start_pos, offset)（inference/model.py:L268-L276）——直接按 (start_pos + 1) // ratio 生成等距位置索引，不走 Indexer 评分。等价的语义是：HCA 假设"重度压缩后的所有位置都重要"，按压缩段的物理位置均匀采样 k=512 个。

(3.6) 稀疏 Attention：

o_t = softmax(q_t · ĉ_K^⊤ / √d) · ĉ_V    # 仅在选中的 k=512 个位置

复杂度分析：

• HCA：O(T²/m' + T·k) = O(T²/128 + T × 512)
• 当 T = 1M：(10¹²/128) + (1.5×10⁹) ≈ 7.8×10⁹ + 1.5×10⁹ ≈ 9.3×10⁹——HCA 的 T² 项几乎消失，主导项是 T·k
• 与 CSA 比较：CSA 的 T²/4 = 2.5×10¹¹，HCA 的 T²/128 = 7.8×10⁹，HCA 在 T² 项上比 CSA 又省 32 倍

Cache 节省：

• HCA 存 ĉ_K, ĉ_V，shape [T/128, d] = [8,192, 4096]
• 比 CSA 的 [T/4, d] 又省 32 倍
• 加上 128-token 滑窗：每层 cache = 128 + 1,048,576/128 = 128 + 8,192 = 8,320 个向量

CSA vs HCA 对比：

HCA 的代价是信息损失更大——每 128 个 token 才保留 1 个压缩向量，必然丢失细节。这也是 V4 把 CSA 和 HCA 逐层交替的原因：让信息保留能力强（m=4）的层和推理计算便宜（m=128）的层互补。下一节展开配合策略。

3.4 CSA + HCA 逐层交替（Hybrid 配合）

V4 实际走的是逐层交替模式——不是分头（query-aware routing），也不是动态（根据输入决定），而是在建模阶段就以 compress_ratios 列表的形式固定下来。V4-Flash 的 compress_ratios（来自 config.json）共 44 项（占位 1 项），前 43 项实际决定每层走哪条分支：

compress_ratios = [
  0,    # Layer 0   纯滑窗 (window=128)
  0,    # Layer 1   纯滑窗 (window=128)
  4,    # Layer 2   CSA (m=4)
  128,  # Layer 3   HCA (m=128)
  4,    # Layer 4   CSA
  128,  # Layer 5   HCA
  ...                              # 4, 128 交替直到 40
  4,    # Layer 40  CSA
  128,  # Layer 41  HCA
  4,    # Layer 42  CSA
  0,    # 占位 (不读)
]

逐层模式总结：

• 前 2 层（Layer 0 与 Layer 1）走纯滑窗（compress_ratios=0），完全不上压缩。原因：浅层需要 raw token 信号建模，让 attention 直接观察原始 token
• Layer 1 走纯滑窗（compress_ratios[1]=0）——是 V4 显式配置（不是异常）：让"非压缩段"覆盖 token 嵌入的早期阶段
• 21 个 CSA 层 + 20 个 HCA 层 + 2 个纯滑窗层 = 43 层（Layer 0/1 滑窗 + Layer 2-42 共 41 层交替 CSA/HCA）

(3.7) 逐层分支：设第 i 层的 attention 行为：

layer_i attention =
  if compress_ratios[i] == 0:    pure sliding window (w=128)
  elif compress_ratios[i] == 4:  CSA branch (m=4, k=512, with Indexer)
  elif compress_ratios[i] == 128: HCA branch (m=128, k=512, position top-k)
  # all branches concat with window_size=128 prefix

共享 MQA-KV + grouped low-rank O 投影：

• MQA-KV：所有 64 个 Q 头共享 1 组 KV（num_key_value_heads=1，head_dim=512），这是 V3 系列 MLA 的核心省 cache 设计
• grouped low-rank O 投影：64 个头分成 8 组（o_groups=8），每组共享一个 o_lora_rank=1024 的低秩矩阵，先 wo_a 把 n_heads × head_dim / n_groups = 4096 维压到 1024 维，再 wo_b 升回 d=4096。这一拆分是 V3 沿用设计，省 cache + 提速
• window_size=128 始终拼接：无论本层是 CSA、HCA 还是纯滑窗，前 128 个位置总是直接走原始 K, V（不上压缩），保证最近的 128 个 token 永远是 raw attention

逐层交替的工程实现：

• 同一 Attention 类（inference/model.py:L436-L543），仅在 __init__ 中按 compress_ratios[layer_id] 选择性创建 Compressor 和 Indexer
  • ratio=0 → 既无 Compressor 也无 Indexer，纯 kv_cache[:bsz, :128]
  • ratio=4 → 创建 Compressor（overlap=True）+ Indexer（Hadamard + FP4）
  • ratio=128 → 创建 Compressor（overlap=False），不创建 Indexer
• forward（inference/model.py:L508-L514）：如果 compress_ratio > 0，计算 compress_topk_idxs（CSA 走 indexer(x, qr, start_pos, offset)，HCA 走 get_compress_topk_idxs(...)），并拼接到 window_topk_idxs 后面

为什么交替而非全部 CSA？

• 全部 CSA：FLOPs 省 4 倍但 cache 仍占 25% of V3.2
• 全部 HCA：cache 省 128 倍但 m=128 损失太多细节，质量下降
• 交替：CSA 隔层"补细节"，HCA 隔层"省算力"——既保质又省 FLOPs/cache

源码入口（完整片段见 §3.6）：

• inference/model.py:L453 — self.compress_ratio = args.compress_ratios[layer_id]
• inference/model.py:L466-L471 — 分支创建逻辑
• inference/model.py:L508-L514 — topk 拼接

下一节给出 V4 官方报告的 FLOPs / KV cache 节省实测数字。

3.5 FLOPs / KV cache 节省分析

V4 技术报告 §1 给出了端到端实测数字（1M 上下文、单 token 推理视角）：

V4-Flash 比 V4-Pro 更激进——Flash 用 m=128 比例更高、hc_mult=4 通道复用更彻底、FP4 专家省 cache。7% / 10% 是 V4-Flash 在 1M 上下文下的目标数字（v0.1 草稿，工程端实测未公开）。

FLOPs 节省的理论估算（自上而下）：

1. 标准 attention（单层单 token）：2 · T · d = 2 × T × 4096
2. CSA：2 · T · (T/m) · d / T + T · k · d = 2 · T · d / m + T · k · d ≈ 2·T·1024 + T·2097152 ≈ T² · 2048 + T · 2.1×10⁶（当 T 主导时简化为 T²/m 项）
3. HCA：2 · T · d / m' + T · k · d ≈ T² · 32 + T · 2.1×10⁶（同样 T 主导时简化为 T²/m' 项）
4. 配合（50% CSA + 50% HCA）：(0.5·T²·1024 + 0.5·T²·32) + T·k·d·2 ≈ 528·T² + 4.2×10⁶·T
   • T=1M 时：528·10¹² + 4.2×10¹² ≈ 5.3×10¹⁴（单层）
   • vs 标准 attention 单层 2·10¹²·4096 = 8.2×10¹⁵
   • 节省比 = 5.3×10¹⁴ / 8.2×10¹⁵ ≈ 6.5%——比 V3.2 还要激进

关键说明：上面的"50% CSA + 50% HCA"是粗略近似。V4-Flash 实际是 21 层 CSA + 20 层 HCA + 2 层纯滑窗的精确配置。V4 官方 10% 数字包含了 ① 滑窗层的全 attention、② Indexer 计算开销、③ YaRN 频率内插、④ grouped low-rank O 投影等多个因子。

KV cache 节省的理论估算：

以上为 KV cache 理论下界，不含 routed expert 权重（FP4 占 ~37 GB）、激活 checkpoint 等运行时开销。V4-Flash 实测 ~17 GB（见 fig-3.5 V4-Flash 柱），比 ~15 GB 略大是因为 MQA-KV 共享、FP4 专家 cache 等细节开销。

为什么 V4-Flash 实测 ~7% 不只是 25%？ 以下几个工程优化叠加：

1. hc_mult=4 残差流复用（CH5）：mHC 让 4 份残差流共享同一份 cache，等价于 cache × 1/4
2. FP4 专家量化（CH7）：routed expert 权重从 FP8 进一步压到 FP4，cache 占用减半
3. MTL 共享 MQA-KV：所有层共享同一组 KV 投影（wkv 一份参数），进一步去重
4. 滑动窗口的物理 cache 仅 128 项：不需要 T × d 那么多 cache

柱状图直观对比（V3.2 = 100%）：

柱状图对 V3.2 MLA / V4 CSA only（理论 4×） / V4 HCA only（理论 32×） / V4 CSA+HCA（实测 10×）做对比；左 Y 轴是 FLOPs/token（log scale），右 Y 轴是 KV cache size（GB）。所有数字基于 V4 技术报告 §1。

结论：V4-Flash 通过"逐层交替 + 重度压缩 + 多项工程优化"三档配合，在 1M 上下文下把单 token FLOPs 压到 V3.2 的 ~10%、KV cache 压到 ~7%——这是 1M 上下文工程化的关键。下一节把上述所有公式映射到实际源码。

3.6 源码映射：CSA + HCA

仓库：deepseek-ai/DeepSeek-V4-Flash
关键文件：inference/model.py、inference/kernel.py

注：行号引用以 inference/model.py 为真源；code-snippets/ 下的 file 段以该 file 自身行号为准。

源码片段 1：Compressor（共享给 CSA / HCA）

对应公式：(3.1) Compress_K + (3.4) HeavyCompress

位置：inference/model.py L279-L378，保存于 code-snippets/compressor.py

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

class Compressor(nn.Module):
    """Compresses KV cache via learned gated pooling over `compress_ratio` consecutive tokens.
    When overlap=True (ratio==4), uses overlapping windows for smoother compression boundaries."""

    def __init__(self, args: ModelArgs, compress_ratio: int = 4, head_dim: int = 512, rotate: bool = False):
        super().__init__()
        self.dim = args.dim
        self.head_dim = head_dim
        self.rope_head_dim = args.rope_head_dim
        self.nope_head_dim = head_dim - args.rope_head_dim
        self.compress_ratio = compress_ratio
        self.overlap = compress_ratio == 4       # <-- overlap only on CSA
        self.rotate = rotate
        coff = 1 + self.overlap

        self.ape = nn.Parameter(torch.empty(compress_ratio, coff * self.head_dim, dtype=torch.float32))
        # wkv and wgate in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for convenient.
        # When overlap, the first half of dims is for overlapping compression, second half for normal.
        self.wkv = Linear(self.dim, coff * self.head_dim, dtype=torch.float32)
        self.wgate = Linear(self.dim, coff * self.head_dim, dtype=torch.float32)
        self.norm = RMSNorm(self.head_dim, args.norm_eps)
        self.kv_cache: torch.Tensor = None  # assigned lazily from Attention.kv_cache
        # State buffers for decode-phase incremental compression.
        # With overlap: state[:, :ratio] = overlapping window, state[:, ratio:] = current window.
        self.register_buffer("kv_state", torch.zeros(args.max_batch_size, coff * compress_ratio, coff * self.head_dim, dtype=torch.float32), persistent=False)
        self.register_buffer("score_state", torch.full((args.max_batch_size, coff * compress_ratio, coff * self.head_dim), float("-inf"), dtype=torch.float32), persistent=False)
        self.freqs_cis: torch.Tensor = None

    def overlap_transform(self, tensor: torch.Tensor, value=0):
        # tensor: [b,s,r,2d]
        b, s, _, _ = tensor.size()
        ratio, d = self.compress_ratio, self.head_dim
        new_tensor = tensor.new_full((b, s, 2 * ratio, d), value)
        new_tensor[:, :, ratio:] = tensor[:, :, :, d:]
        new_tensor[:, 1:, :ratio] = tensor[:, :-1, :, :d]
        return new_tensor

    def forward(self, x: torch.Tensor, start_pos: int):
        assert self.kv_cache is not None
        bsz, seqlen, _ = x.size()
        ratio, overlap, d, rd = self.compress_ratio, self.overlap, self.head_dim, self.rope_head_dim
        dtype = x.dtype
        # compression need fp32
        x = x.float()
        kv = self.wkv(x)
        score = self.wgate(x)
        if start_pos == 0:
            should_compress = seqlen >= ratio
            remainder = seqlen % ratio
            cutoff = seqlen - remainder
            offset = ratio if overlap else 0
            if overlap and cutoff >= ratio:
                self.kv_state[:bsz, :ratio] = kv[:, cutoff-ratio : cutoff]
                self.score_state[:bsz, :ratio] = score[:, cutoff-ratio : cutoff] + self.ape
            if remainder > 0:
                kv, self.kv_state[:bsz, offset : offset+remainder] = kv.split([cutoff, remainder], dim=1)
                self.score_state[:bsz, offset : offset+remainder] = score[:, cutoff:] + self.ape[:remainder]
                score = score[:, :cutoff]
            kv = kv.unflatten(1, (-1, ratio))
            score = score.unflatten(1, (-1, ratio)) + self.ape
            if overlap:
                kv = self.overlap_transform(kv, 0)
                score = self.overlap_transform(score, float("-inf"))
            kv = (kv * score.softmax(dim=2)).sum(dim=2)
        else:
            should_compress = (start_pos + 1) % self.compress_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:
                    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:]
                    self.score_state[:bsz, :ratio] = self.score_state[:bsz, ratio:]
            else:
                self.kv_state[:bsz, start_pos % ratio] = kv.squeeze(1)
                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)
        if not should_compress:
            return
        kv = self.norm(kv.to(dtype))
        if start_pos == 0:
            freqs_cis = self.freqs_cis[:cutoff:ratio]
        else:
            freqs_cis = self.freqs_cis[start_pos + 1 - self.compress_ratio].unsqueeze(0)
        apply_rotary_emb(kv[..., -rd:], freqs_cis)
        if self.rotate:
            kv = rotate_activation(kv)
            fp4_act_quant(kv, fp4_block_size, True)
        else:
            act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True)
        if start_pos == 0:
            self.kv_cache[:bsz, :seqlen // ratio] = kv
        else:
            self.kv_cache[:bsz, start_pos // ratio] = kv.squeeze(1)
        return kv

逐段说明：

• L290 self.overlap = compress_ratio == 4 —— 关键：仅 CSA 启用重叠，HCA 不开
• L294-L298 4 个核心参数：ape（绝对位置编码，加到 softmax 分数上）、wkv（K/V 投影）、wgate（gating 投影）、norm（RMSNorm）
• L316-L378 forward：分 start_pos == 0（prefill）和 start_pos > 0（decode）两条路径
  • L342 kv = (kv * score.softmax(dim=2)).sum(dim=2) —— softmax-pool 压缩的核心：对 ratio 个 token 算 softmax 权重再加权求和
  • L362 apply_rotary_emb(kv[..., -rd:], freqs_cis) —— 压缩后施 RoPE
  • L368-L372 if self.rotate（Indexer 内的 Compressor）：FP4 模拟；else（Attention 共用的 Compressor）：FP8 模拟

为可读性省略：类外辅助函数 apply_rotary_emb、act_quant、fp4_act_quant、rotate_activation 的具体实现（详见 inference/kernel.py）。

源码片段 2：Indexer（CSA 专属）

对应公式：(3.2) Indexer 评分 + Index_score 计算

位置：inference/model.py L380-L434，保存于 code-snippets/indexer.py

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

class Indexer(torch.nn.Module):
    """Selects top-k compressed KV positions for sparse attention via learned scoring.
    Has its own Compressor (with Hadamard rotation) to build compressed KV for scoring."""

    def __init__(self, args: ModelArgs, compress_ratio: int = 4):
        super().__init__()
        self.dim = args.dim
        self.n_heads = args.index_n_heads
        self.n_local_heads = args.index_n_heads // world_size
        self.head_dim = args.index_head_dim
        self.rope_head_dim = args.rope_head_dim
        self.index_topk = args.index_topk
        self.q_lora_rank = args.q_lora_rank
        self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim)
        self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16)
        self.softmax_scale = self.head_dim ** -0.5
        self.compress_ratio = compress_ratio

        self.compressor = Compressor(args, compress_ratio, self.head_dim, True)
        self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len // compress_ratio, self.head_dim), persistent=False)
        self.freqs_cis = None

    def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, offset: int):
        bsz, seqlen, _ = x.size()
        freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]
        ratio = self.compress_ratio
        rd = self.rope_head_dim
        end_pos = start_pos + seqlen
        if self.compressor.kv_cache is None:
            self.compressor.kv_cache = self.kv_cache
            self.compressor.freqs_cis = self.freqs_cis
        q = self.wq_b(qr)
        q = q.unflatten(-1, (self.n_local_heads, self.head_dim))
        apply_rotary_emb(q[..., -rd:], freqs_cis)
        q = rotate_activation(q)
        # use fp4 simulation for q and kv in indexer
        fp4_act_quant(q, fp4_block_size, True)
        self.compressor(x, start_pos)
        weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)
        # We performed QAT here, kv could also use fp8 format, though current implementation uses bf16
        index_score = torch.einsum("bshd,btd->bsht", q, self.kv_cache[:bsz, :end_pos // ratio])
        index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)
        if world_size > 1:
            dist.all_reduce(index_score)
        if start_pos == 0:
            mask = torch.arange(seqlen // ratio).repeat(seqlen, 1) >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio
            index_score += torch.where(mask, float("-inf"), 0)
        topk_idxs = index_score.topk(min(self.index_topk, end_pos // ratio), dim=-1)[1]
        if start_pos == 0:
            mask = topk_idxs >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio
            topk_idxs = torch.where(mask, -1, topk_idxs + offset)
        else:
            topk_idxs += offset
        return topk_idxs

逐段说明：

• L398 self.compressor = Compressor(args, compress_ratio, self.head_dim, True) —— Indexer 内的 Compressor 与 Attention 共用的 Compressor 是两个不同实例，rotate=True 触发 Hadamard + FP4 路径
• L411-L416 q 投影：先 wq_b（共用 Attention 的 wq_b），再 RoPE + Hadamard + FP4
• L418 weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5) —— 每个 query 一个权重 w_t，缩放因子是 1/√d × 1/√h
• L420 index_score = einsum("bshd,btd->bsht", q, kv_cache) —— 关键：query × key 评分
• L421 index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2) —— 公式 (3.2) 完整实现：relu(q·k^T) × w_t → 求和
• L427 topk_idxs = index_score.topk(min(self.index_topk, end_pos // ratio), dim=-1)[1] —— 选 top-512

为可读性省略：辅助函数 ColumnParallelLinear、apply_rotary_emb、fp4_act_quant 的实现细节。

源码片段 3：Attention（CSA + HCA 整合）

对应公式：(3.3) Sparse Attention + (3.6) Position top-k + (3.7) Hybrid combine

位置：inference/model.py L436-L543，保存于 code-snippets/attention.py

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108

class Attention(nn.Module):
    """Multi-head Latent Attention (MLA) with sliding window + optional KV compression.
    Uses low-rank Q projection (wq_a -> q_norm -> wq_b) and grouped low-rank O projection."""
    def __init__(self, layer_id: int, args: ModelArgs):
        super().__init__()
        self.layer_id = layer_id
        self.dim = args.dim
        self.n_heads = args.n_heads
        self.n_local_heads = args.n_heads // world_size
        self.q_lora_rank = args.q_lora_rank
        self.o_lora_rank = args.o_lora_rank
        self.head_dim = args.head_dim
        self.rope_head_dim = args.rope_head_dim
        self.nope_head_dim = self.head_dim - self.rope_head_dim
        self.n_groups = args.o_groups
        self.n_local_groups = self.n_groups // world_size
        self.window_size = args.window_size
        self.compress_ratio = args.compress_ratios[layer_id]   # <-- 逐层选择
        self.eps = args.norm_eps

        self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32))
        self.wq_a = Linear(self.dim, self.q_lora_rank)
        self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
        self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim)
        self.wkv = Linear(self.dim, self.head_dim)
        self.kv_norm = RMSNorm(self.head_dim, self.eps)
        self.wo_a = ColumnParallelLinear(self.n_heads * self.head_dim // self.n_groups, self.n_groups * args.o_lora_rank, dtype=torch.bfloat16)
        self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank, self.dim)
        self.softmax_scale = self.head_dim ** -0.5

        if self.compress_ratio:
            self.compressor = Compressor(args, self.compress_ratio, self.head_dim)
            if self.compress_ratio == 4:
                self.indexer = Indexer(args, self.compress_ratio)   # <-- CSA 才有
            else:
                self.indexer = None                                 # <-- HCA 没有

        kv_cache_size = args.window_size + (args.max_seq_len // self.compress_ratio if self.compress_ratio else 0)
        self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, kv_cache_size, self.head_dim), persistent=False)
        if self.compress_ratio:
            original_seq_len, rope_theta = args.original_seq_len, args.compress_rope_theta
        else:
            # disable YaRN and use base rope_theta in pure sliding-window attention
            original_seq_len, rope_theta = 0, args.rope_theta
        freqs_cis = precompute_freqs_cis(self.rope_head_dim, args.max_seq_len, original_seq_len,
                                         rope_theta, args.rope_factor, args.beta_fast, args.beta_slow)
        self.register_buffer("freqs_cis", freqs_cis, persistent=False)

    def forward(self, x: torch.Tensor, start_pos: int):
        bsz, seqlen, _ = x.size()
        freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]
        win = self.window_size
        ratio = self.compress_ratio
        rd = self.rope_head_dim
        if self.compress_ratio and self.compressor.kv_cache is None:
            self.compressor.kv_cache = self.kv_cache[:, win:]
            self.compressor.freqs_cis = self.freqs_cis
            if self.indexer is not None:
                self.indexer.freqs_cis = self.freqs_cis
        # q
        qr = q = self.q_norm(self.wq_a(x))
        q = self.wq_b(q).unflatten(-1, (self.n_local_heads, self.head_dim))
        q *= torch.rsqrt(q.square().mean(-1, keepdim=True) + self.eps)
        apply_rotary_emb(q[..., -rd:], freqs_cis)

        # win kv & topk_idxs
        kv = self.wkv(x)
        kv = self.kv_norm(kv)
        apply_rotary_emb(kv[..., -rd:], freqs_cis)
        # FP8-simulate non-rope dims to match QAT; rope dims stay bf16 for positional precision
        act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True)
        topk_idxs = get_window_topk_idxs(win, bsz, seqlen, start_pos)
        if self.compress_ratio:
            offset = kv.size(1) if start_pos == 0 else win
            if self.indexer is not None:                          # <-- CSA
                compress_topk_idxs = self.indexer(x, qr, start_pos, offset)
            else:                                                 # <-- HCA
                compress_topk_idxs = get_compress_topk_idxs(ratio, bsz, seqlen, start_pos, offset)
            topk_idxs = torch.cat([topk_idxs, compress_topk_idxs], dim=-1)   # <-- window + compress 拼接
        topk_idxs = topk_idxs.int()

        # compress kv & attn
        if start_pos == 0:
            if seqlen <= win:
                self.kv_cache[:bsz, :seqlen] = kv
            else:
                cutoff = seqlen % win
                self.kv_cache[:bsz, cutoff: win], self.kv_cache[:bsz, :cutoff] = kv[:, -win:].split([win - cutoff, cutoff], dim=1)
            if self.compress_ratio:
                if (kv_compress := self.compressor(x, start_pos)) is not None:
                    kv = torch.cat([kv, kv_compress], dim=1)
            # We performed QAT here, kv could also use fp8 format, though current implementation uses bf16
            o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)
        else:
            self.kv_cache[:bsz, start_pos % win] = kv.squeeze(1)
            if self.compress_ratio:
                self.compressor(x, start_pos)
            o = sparse_attn(q, self.kv_cache[:bsz], self.attn_sink, topk_idxs, self.softmax_scale)
        apply_rotary_emb(o[..., -rd:], freqs_cis, True)

        # o
        o = o.view(bsz, seqlen, self.n_local_groups, -1)
        wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
        # NOTE: wo_a is FP8 in checkpoint; could do FP8 einsum here for better perf,
        # but using BF16 for simplicity.
        o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
        x = self.wo_b(o.flatten(2))
        return x

逐段说明：

• L453 self.compress_ratio = args.compress_ratios[layer_id] —— 公式 (3.7) 的层选择开关
• L466-L471 关键分支：if self.compress_ratio: ... if self.compress_ratio == 4: self.indexer = ... else: self.indexer = None —— CSA 创建 Indexer，HCA 不创建
• L473 kv_cache_size = window_size + max_seq_len // compress_ratio —— 实际 cache 分配：滑窗 + 压缩段
• L475-L479 YaRN 仅在 compress_ratio > 0 时启用，滑窗层用基础 rope_theta
• L502-L506 滑窗段的 K, V 计算 + RoPE + FP8 模拟
• L507 topk_idxs = get_window_topk_idxs(win, bsz, seqlen, start_pos) —— 128-token 滑窗位置
• L508-L514 公式 (3.7) 完整实现：
  • ratio=4（CSA）：compress_topk_idxs = self.indexer(x, qr, start_pos, offset) —— 走 Indexer
  • ratio=128（HCA）：compress_topk_idxs = get_compress_topk_idxs(ratio, bsz, seqlen, start_pos, offset) —— 走位置 top-k
  • 两者都拼到 topk_idxs 后面
• L518-L533 prefill vs decode 两条路径调用 sparse_attn kernel
• L537-L542 grouped low-rank O 投影：8 组头共享 wo_a 矩阵

为可读性省略：get_window_topk_idxs、get_compress_topk_idxs、sparse_attn、apply_rotary_emb、act_quant 等辅助函数。

源码片段 4：sparse_attn_kernel（TileLang 实现）

对应公式：(3.3) Sparse Attention 的 GPU kernel

位置：inference/kernel.py L276-L368，保存于 code-snippets/sparse_attn_kernel.py

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

@tilelang.jit(pass_configs=pass_configs)
def sparse_attn_kernel(h: int, d: int, scale=None):
    """Sparse multi-head attention via index gathering + online softmax (FlashAttention-style).
    For each (batch, seq_pos), gathers top-k KV positions by index, computes attention
    with numerically stable running max/sum, and includes a learnable attn_sink bias."""
    b = T.symbolic("b")
    m = T.symbolic("m")
    n = T.symbolic("n")
    topk = T.symbolic("topk")
    if scale is None:
        scale = (1.0 / d) ** 0.5

    num_stages = 2
    threads = 256
    block = 64
    num_blocks = tilelang.cdiv(topk, block)

    @T.prim_func
    def sparse_attn_kernel_(
        q: T.Tensor[(b, m, h, d), BF16],
        kv: T.Tensor[(b, n, d), BF16],
        o: T.Tensor[(b, m, h, d), BF16],
        attn_sink: T.Tensor[(h,), FP32],
        topk_idxs: T.Tensor[(b, m, topk), INT32],
    ):
        with T.Kernel(m, b, threads=threads) as (bx, by):
            q_shared = T.alloc_shared((h, d), BF16)
            kv_shared = T.alloc_shared((block, d), BF16)
            o_shared = T.alloc_shared((h, d), BF16)
            acc_s_cast = T.alloc_shared((h, block), BF16)

            idxs = T.alloc_fragment(block, INT32)
            acc_s = T.alloc_fragment((h, block), FP32)
            acc_o = T.alloc_fragment((h, d), FP32)
            scores_max = T.alloc_fragment(h, FP32)
            scores_max_prev = T.alloc_fragment(h, FP32)
            scores_scale = T.alloc_fragment(h, FP32)
            scores_sum = T.alloc_fragment(h, FP32)
            sum_exp = T.alloc_fragment(h, FP32)

            T.clear(acc_o)
            T.clear(sum_exp)
            T.fill(scores_max, -T.infinity(FP32))
            T.copy(q[by, bx, :, :], q_shared)

            for t in T.Pipelined(num_blocks, num_stages=num_stages):
                for i in T.Parallel(block):
                    idxs[i] = T.if_then_else(t * block + i < topk, topk_idxs[by, bx, t * block + i], -1)
                for i, j in T.Parallel(block, d):
                    kv_shared[i, j] = T.if_then_else(idxs[i] != -1, kv[by, idxs[i], j], 0)
                for i, j in T.Parallel(h, block):
                    acc_s[i, j] = T.if_then_else(idxs[j] != -1, 0, -T.infinity(FP32))
                T.gemm(q_shared, kv_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
                for i, j in T.Parallel(h, block):
                    acc_s[i, j] *= scale
                T.copy(scores_max, scores_max_prev)
                T.reduce_max(acc_s, scores_max, dim=1, clear=False)
                for i in T.Parallel(h):
                    scores_scale[i] = T.exp(scores_max_prev[i] - scores_max[i])
                for i, j in T.Parallel(h, block):
                    acc_s[i, j] = T.exp(acc_s[i, j] - scores_max[i])
                T.reduce_sum(acc_s, scores_sum, dim=1)
                for i in T.Parallel(h):
                    sum_exp[i] = sum_exp[i] * scores_scale[i] + scores_sum[i]
                T.copy(acc_s, acc_s_cast)
                for i, j in T.Parallel(h, d):
                    acc_o[i, j] *= scores_scale[i]
                T.gemm(acc_s_cast, kv_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)

            for i in T.Parallel(h):
                sum_exp[i] += T.exp(attn_sink[i] - scores_max[i])
            for i, j in T.Parallel(h, d):
                acc_o[i, j] /= sum_exp[i]
            T.copy(acc_o, o_shared)
            T.copy(o_shared, o[by, bx, :, :])

    return sparse_attn_kernel_

逐段说明：

• L301 with T.Kernel(m, b, threads=threads) as (bx, by) —— 启动 grid：每个 (batch, query_position) 一个 block
• L321-L323 topk_idxs[by, bx, t * block + i] —— 关键：按 topk 索引 gather kv，而非按连续位置
• L324-L325 kv[by, idxs[i], j] —— 公式 (3.3) 稀疏 attention 的核心：从 [b, n, d] 张量 gather 出 top-k 个位置
• L328 T.gemm(q_shared, kv_shared, acc_s, transpose_B=True) —— 算 Q · K^⊤
• **L331-L339FlashAttention 风格的 online softmax：跟踪scores_max、scores_scale、sum_exp，无需完整物化 [h, topk]` 矩阵
• L343 T.gemm(acc_s_cast, kv_shared, acc_o) —— 算 softmax · V
• L346 sum_exp[i] += T.exp(attn_sink[i] - scores_max[i]) —— attn_sink 是 V3.2 沿用的设计：每个 head 一个标量 bias，加到 softmax 分母上，让 token 概率不至于衰减到 0
• **L347-L350最终归一化 + 写回o`

为可读性省略：TileLang 的 @T.prim_func 装饰器、tilelang.jit 编译配置（pass_configs）等基础设施。

形式化 ↔ 代码 映射总表

章节小结：CH3 把 V4-Flash 注意力机制的"灾难 → CSA → HCA → 配合 → 实测 → 源码"完整走了一遍。关键创新在三个层面：(1) 用 Compressor 把 K, V 沿时间维做有学习权重的 softmax-pool 压缩；(2) CSA 用 Indexer 评分选 top-k，HCA 直接用位置 top-k 省 Indexer；(3) 逐层交替 + 滑窗 + MQA-KV + grouped low-rank O 投影四档配合，让 1M 上下文下的 FLOPs / cache 同时降到 V3.2 的 10% / 7% 量级。

3.7 算子级拆解：Compressor / Indexer / sparse_attn_kernel

本节与 3.6 的区别：3.6 给的是"完整源码 + 逐行注释"——侧重把公式 (3.1)-(3.7) 跟具体行号对位；3.7 给的是"为什么这样写代码"——把 3 个最关键的算子（Compressor.forward / Indexer.forward / sparse_attn_kernel）拆到算子级，解释 (a) 边界条件为什么这样分支、(b) 数值缩放因子从哪儿来、(c) kernel 内部 online softmax 的状态机怎么循环。

3.7.1 Compressor：overlap 边界处理

Compressor.forward 是 V4-Flash 整个时间压缩的核心算子。它处理三种 case：(1) prefill 全段整批压缩；(2) decode 单 token 流入；(3) overlap 模式下 2 段拼接。下面的代码片段是 §3.6 源码片段 1 的"去壳版"——只保留三种 case 的核心分支，删掉了 act_quant / RoPE 等与边界无关的细节：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

def forward(self, x: torch.Tensor, start_pos: int):
    bsz, seqlen, _ = x.size()
    ratio, overlap, d = self.compress_ratio, self.overlap, self.head_dim
    kv   = self.wkv(x)         # [b, s, d]
    score = self.wgate(x)      # [b, s, d]，softmax-pool 的 logit

    if start_pos == 0:
        # (A) PREFILL: 把整段 seqlen 切成 ratio-token 组, 整段 softmax-pool
        cutoff  = seqlen - seqlen % ratio
        kv      = kv[:, :cutoff].unflatten(1, (-1, ratio))      # [b, n_group, ratio, d]
        score   = score[:, :cutoff].unflatten(1, (-1, ratio)) + self.ape
        if overlap:
            kv    = self.overlap_transform(kv, 0)                # [b, n_group, 2*ratio, d]
            score = self.overlap_transform(score, float("-inf")) # 边界外位置用 -inf 填
        kv = (kv * score.softmax(dim=2)).sum(dim=2)              # [b, n_group, d] (overlap 2d)
    else:
        # (B) DECODE: 边走边压缩, 用 kv_state (ratio x d) + score_state (ratio x d) 维护
        slot = start_pos % ratio
        score = score + self.ape[slot]
        if overlap:
            self.kv_state[:bsz, ratio + slot]   = kv.squeeze(1)   # 写入右半窗口
            self.score_state[:bsz, ratio + slot] = score.squeeze(1)
            if (start_pos + 1) % ratio == 0:                     # 凑齐 ratio 个 token
                kv_state   = torch.cat([self.kv_state[:bsz, :ratio, :d],
                                        self.kv_state[:bsz, ratio:, d:]], dim=1)  # 沿 d 拼
                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:]
                self.score_state[:bsz, :ratio] = self.score_state[:bsz, ratio:]
        else:
            self.kv_state[:bsz, slot]   = kv.squeeze(1)           # 写入 slot
            self.score_state[:bsz, slot] = score.squeeze(1)
            if (start_pos + 1) % ratio == 0:
                kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True)

四个关键设计点：

(1) start_pos == 0（prefill 路径）：cutoff = seqlen - seqlen % ratio 把整段切成 n_group = seqlen // ratio 个 ratio-token 段；unflatten(1, (-1, ratio)) 把 [b, s, d] reshape 成 [b, n_group, ratio, d]，再在 dim=2 上做 softmax + sum——这是有学习权重的 softmax-pool 公式 (3.1) 的向量化实现，比 for-loop 跑得快。self.ape 是 per-position 的 additive positional encoding（[ratio, 2d]），加到 score 上后让不同位置有不等的"被选中概率"。

(2) start_pos > 0（decode 路径）：Compressor 维护两个 buffer：kv_state [B, 2*ratio, 2*d] 和 score_state [B, 2*ratio, 2*d]（HCA 时第二个维度退化为 1*ratio）。每来一个新 token，先把它写进 kv_state[bsz, slot] 和 score_state[bsz, slot]，再判断 should_compress = (start_pos+1) % ratio == 0——凑齐 ratio 个 token 就做一次 softmax-pool。这一段逻辑等价于一个"循环长度为 ratio 的 FIFO 累加器"，复杂度 O(1) per token。

(3) overlap = (compress_ratio == 4)：只有 CSA（m=4）才会触发 overlap，HCA（m=128）overlap=False。overlap 模式下 coff = 1 + overlap = 2，kv_state 容量扩到 2×ratio=8。每凑齐 4 个 token 后用下面这段（公式 3.8）一次性出 2 段压缩向量：

其中 cat([kv_state[:, :4, :d], kv_state[:, 4:, d:]], dim=1) 是把左半窗口的"原始 d 维"和右半窗口的"另 d 维"沿 d 维拼起来——这是 overlap_transform 在 decode 侧的逆操作（prefill 用 overlap_transform 展开 [b, s, r, 2d] → [b, s, 2r, d]，decode 用 cat 沿 d 拼回）。关键观察：2m=8 个 token 中，t1..t4 全部进了 group A 的 softmax，t5..t8 全部进了 group B 的 softmax——边界处不重叠；overlap 的实际效果是让"通道维"加倍（每个 token 给出 2d 维表示，前 d 维供前组、后 d 维供后组），从而让 group A / B 的 softmax-pool 各自用上全部 4 个 token 的信息而不是被切碎。

(4) HCA 不开 overlap：当 compress_ratio=128 时 overlap=False，每 128 个 token 才压缩 1 次，边界处不重叠。kv_state 容量 1*ratio=128，should_compress = (start_pos+1) % 128 == 0。128 的 ratio 已经足够稀疏（1M / 128 ≈ 8192 个压缩位置），再用 overlap 反而会让 kv_state 显存翻倍——性价比太低。所以 V4 的设计是"只在 ratio 极小（m=4）时才花这个开销"。

3.7.2 Indexer：w_t + softmax_scale 的两段缩放

Indexer.forward 是 CSA 专属（compress_ratio==4 时启用）。它的核心动作只有两行——算 index_score 和 topk——但里面藏着 4 个微妙设计。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

def forward(self, x, qr, start_pos, offset):
    # 1. q 端: wq_b + RoPE + Hadamard + FP4 模拟 (QAT 准备)
    q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim))  # [b, s, h=64, d=128]
    apply_rotary_emb(q[..., -rd:], freqs_cis)
    q = rotate_activation(q)                    # Hadamard 旋转
    fp4_act_quant(q, fp4_block_size, True)       # 模拟 FP4 (QAT)
    self.compressor(x, start_pos)               # Indexer 自带 Compressor (rotate=True)
    # 2. 核心: per-head 权重 w_t
    weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)
    # 3. index_score: q 与 (T/4) 个压缩 key 的内积
    index_score = torch.einsum("bshd,btd->bsht", q, self.kv_cache[:bsz, :end_pos // ratio])
    # 4. ReLU + 权重求和 = 公式 (3.2)
    index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)
    topk_idxs = index_score.topk(min(self.index_topk, end_pos // ratio), dim=-1)[1]
    return topk_idxs

四个细节：

(1) weights_proj 的形状：ColumnParallelLinear(dim, n_heads)——把每个 token 的 d 维投影成 n_heads=64 维的per-head 标量权重。为什么是 per-head 不是 per-token：index_score 是 [bsz, seq, n_heads, n_compressed]，64 头对 T/4 个压缩位置各自打分；w_t 也需要是 per-head 的，让每头有独立的"query 重要性"调节。如果 w_t 是 per-token 标量（shape=[bsz, seq, 1]），所有 64 头会被同一个权重拉伸，损失"头间特异化"的能力。

(2) 1/√d · 1/√h 缩放因子的来源：weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)——self.softmax_scale = head_dim ** -0.5 = d ** -0.5（head_dim=128）。前一半 1/√d 是 attention 的标准缩放（让 q·k^⊤ 数值落在合理范围），后一半 1/√h 是按头数归一化——让 64 头 w_t 的平均权重稳定在 1 附近，避免头数多导致权重爆炸。这两个缩放合起来，让公式 (3.2) 的 relu(q · k^⊤) · w_t 的输出均值稳定在 O(1) 量级，topk 选位置时不会被某一个 head 的偏大权重主导。

(3) index_score = einsum("bshd,btd->bsht", q, kv_cache)：64 head × 128 dim 的 query 与 T/4 个压缩 key 算分数。为什么是 bshd → bsht 而不是 bshd → bshs（self-attention）：Indexer 不算 self-attention，而是把 64 head 的 query 都拿去和压缩 K（kv_cache 是 Indexer 自己的压缩缓存，不是 Attention 那个）做点积——得到的 bsht 含义是"每个 query 位置对每个压缩位置的 64 个 head-specific 分数"。下一步要 sum(dim=2) 把 64 头折成一个标量。

(4) ReLU vs softmax 选哪个：index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)——公式 (3.2) 的完整实现。为什么用 ReLU 不是 softmax：index_score 是"sparse top-k 选择器"，只需要非负权重 + 选择性激活；softmax 会把分数归一化成概率分布（O(topk) 个 exp 计算），而 ReLU 是 O(1) 的逐元素 max(0, x)。对 1M 上下文（T/4 ≈ 262144 个压缩位置、64 头），softmax 会让 indexer 算量增加 1 个 exp / 位置，ReLU 完全省掉这笔开销。质量影响：ReLU 把负分数直接归零，相当于"硬截断"——但 indexer 的目的是"挑出 top-k"，负分数本来就不会进 top-k，截断不损失信息。

(5) Indexer 内的 Compressor vs Attention 共用的 Compressor：

• Indexer 内的：Compressor(args, compress_ratio, head_dim, rotate=True)（indexer.py:L19）— 触发 Hadamard 旋转 + FP4 模拟（fp4_act_quant 在 indexer.py:L37），让 indexer 的输入/输出都对 FP4 量化更鲁棒
• Attention 共用的：Compressor(args, compress_ratio, head_dim, rotate=False)（attention.py:L32）— 走 act_quant (FP8)

为什么需要两套：Indexer 是"sparse 选位置"——只关心哪些压缩位置最相关，对单个压缩向量的数值精度要求低，可以激进量化（FP4）不损失 top-k 准确率；Attention 是"实际算 attention output"——压缩向量要进 softmax(Q·K^⊤/√d)·V 算最终输出，量化要保守（FP8）保质量。这也是 V3.2 → V4 的关键经验：压缩-评分-选择这条链路用 FP4，压缩-实际-计算这条链路用 FP8，两套 Compressor 共享结构但走不同量化路径。

3.7.3 sparse_attn_kernel：TileLang kernel 内部机制

sparse_attn_kernel 是 V4-Flash 实际算 attention output 的算子，编译期把 h=64, d=128, scale=1/√d 常量化、运行期把 b, m, n, topk 做成符号维度。完整代码片段：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

@tilelang.jit
def sparse_attn_kernel(h, d, scale=1.0/d**0.5):
    block = 64; num_stages = 2
    num_blocks = tilelang.cdiv(topk, block)  # 512/64 = 8

    @T.prim_func
    def sparse_attn_kernel_(q, kv, o, attn_sink, topk_idxs):
        with T.Kernel(m, b, threads=256) as (bx, by):
            # 5 个 buffer (shared + fragment)
            q_shared   = T.alloc_shared((h, d), BF16)         # 64 头 × 128 dim query
            kv_shared  = T.alloc_shared((block, d), BF16)      # 64 个压缩 K/V
            acc_s      = T.alloc_fragment((h, block), FP32)   # 当前 block 分数
            acc_o      = T.alloc_fragment((h, d), FP32)       # 当前 head output 累加
            o_shared   = T.alloc_shared((h, d), BF16)         # 最终写回
            # 4 个 online softmax buffer
            scores_max      = T.alloc_fragment(h, FP32)
            scores_max_prev = T.alloc_fragment(h, FP32)
            scores_scale    = T.alloc_fragment(h, FP32)
            scores_sum      = T.alloc_fragment(h, FP32)
            T.fill(scores_max, -T.infinity(FP32))
            T.copy(q[by, bx], q_shared)
            for t in T.Pipelined(num_blocks, num_stages=num_stages):
                # 1. gather top-k 位置 -> kv_shared
                idxs[i] = topk_idxs[by, bx, t*block + i]
                kv_shared[i, j] = kv[by, idxs[i], j]
                # 2. QK^T -> acc_s
                T.gemm(q_shared, kv_shared, acc_s, transpose_B=True)
                acc_s *= scale
                # 3. online softmax: max + rescale + sum
                T.copy(scores_max, scores_max_prev)
                T.reduce_max(acc_s, scores_max, dim=1, clear=False)
                scores_scale = T.exp(scores_max_prev - scores_max)
                acc_s = T.exp(acc_s - scores_max)              # 减去 max 防溢出
                T.reduce_sum(acc_s, scores_sum, dim=1)
                scores_sum = scores_sum_prev * scores_scale + scores_sum
                # 4. PV: acc_o *= scale, acc_o += acc_s @ kv
                acc_o *= scores_scale
                T.gemm(acc_s_cast, kv_shared, acc_o)
            # 5. attn_sink: V3.2 沿用的偏置 (per-head 标量)
            scores_sum += T.exp(attn_sink - scores_max)
            acc_o /= scores_sum
            T.copy(acc_o, o_shared)
            T.copy(o_shared, o[by, bx])
    return sparse_attn_kernel_

五个关键机制：

(1) kernel 入口签名 + 编译期常量化：sparse_attn_kernel(h=64, d=128, scale=1/√d) 把 head 数和 head_dim 作为 Python 参数传进去，在 @T.prim_func 装饰时 TileLang 会把它们编译期常量化——这意味着循环展开、register 分配都可以针对 h*d=8192 元素大小做特化，避免运行期动态形状带来的开销。运行期只有 (b, m, n, topk) 是符号维度。

(2) 5 个 buffer + 4 个 online softmax buffer：

• q_shared: (h, d) bf16 — 64 头 × 128 dim query
• kv_shared: (block, d) bf16 — 64 个压缩 K/V
• acc_s: (h, block) fp32 — 当前 block 的 attention 分数
• acc_o: (h, d) fp32 — 当前 head 的 attention output 累加
• o_shared: (h, d) bf16 — 最终写回
• 4 个 online softmax buffer：scores_max, scores_max_prev, scores_scale, scores_sum（FlashAttention 标准三件套 + prev 备份）

注意 acc_* 都是 fp32（fragment 寄存器）—— 累加用 fp32 避免 bf16 的尾数截断；q_shared / kv_shared 是 bf16（shared memory）—— 输入/输出用 bf16 节约 shared memory 带宽。

(3) online softmax 三件套（重点）：

• scores_max 跟踪历史最大值（init = -inf，每个 block 后 max 更新）
• scores_scale = exp(scores_max_prev - scores_max) 是 rescale 因子（让旧的累加 acc_o / sum_exp 乘上这个因子，相当于在新的 max 基准下重新归一化）
• scores_sum 累加 exp(scores - scores_max)（减 max 是为防 exp 溢出）
• 每进一个新 block：acc_o = acc_o * scores_scale + 新 block 贡献，最后 o = acc_o / scores_sum

(4) block 切分 + pipeline：

• block=64：每 64 个 top-k 位置为 1 个 block
• num_stages=2：双缓冲 pipeline（边算当前 block 边加载下一 block 到 shared memory）
• num_blocks = topk / 64 = 512 / 64 = 8：每 query 8 个 block

T.Pipelined(num_blocks, num_stages=2) 是 TileLang 提供的"软件流水线"原语——编译器会自动插入双缓冲，让 block t+1 的 kv_shared 加载与 block t 的 GEMM 计算重叠。这对稀疏 attention 尤其重要：每个 block 的 kv_shared 都从 global memory 重新 gather（按 topk_idxs 随机索引），没有 locality，必须用 pipeline 隐藏 gather 延迟。

(5) attn_sink 机制（V3.2 沿用）：

1
2
3

attn_sink: T.Tensor[(h,), FP32]    # per-head 标量
...
sum_exp[i] += T.exp(attn_sink[i] - scores_max[i])

attn_sink 是 V3.2 时代就引入的设计——Attention.__init__ L21 里 self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32))，每个 head 一个标量，在 online softmax 末尾加进 sum_exp。语义：让"无 K/V 位置"的"虚拟 token"对所有 query 贡献一个固定偏置（类似 attention 中的"空槽位"），保证 attention 不完全塌缩——如果某些 query 的 top-k 分数都很低（被 indexer 选错了），attn_sink 仍然给一个 fallback 概率，让 output 不会变成 NaN。V4 完整保留 V3.2 的 attn_sink 设计，与 HCA 无关——HCA 没有 indexer 也照样有 attn_sink。

章节小结：3.7 把 3.6 引用过的 3 个核心算子拆到算子级，补充了"为什么这样写代码"的设计动机：

• Compressor：3 个分支（prefill / decode / overlap）+ kv_state/score_state 累加器 + overlap=(ratio==4) 的差异化设计 → 让 m=4 时边界平滑、m=128 时省显存
• Indexer：per-head w_t 让 64 头独立调节 + 1/√d · 1/√h 双缩放稳定数值 + ReLU 省 exp 开销 + Indexer-Compressor（rotate=True FP4）vs Attention-Compressor（rotate=False FP8）两套量化路径
• sparse_attn_kernel：编译期常量化 h,d + 5 buffer + 4 online softmax 状态 + block=64 num_stages=2 num_blocks=8 三档 pipeline + V3.2 沿用的 attn_sink 偏置

下一章展开 MoE 路由。

CH 4. MoE 路由：1 + 256 + 6 拓扑

CH3 解决了"每个 token 怎么挑 512 个历史位置"的问题，CH4 解决"每个 token 怎么挑 6 个专家 + 1 个共享"的问题。V4-Flash 的 MoE 是 256 专家库 + 1 个永远激活的共享专家 + 每 token 激活 6 个 routed 专家的稀疏混合结构（DeepSeekMoE 范式），并新增了 hash routing（前三层）、sqrtsoftplus 评分、aux-loss-free 偏置三大 V3 增量。

4.1 1 + 256 + 6 拓扑结构

V4-Flash 的 MoE 沿用 DeepSeekMoE（Dai et al., 2024）“细粒度专家 + 共享专家隔离"的设计：把稠密 FFN 拆成 E 个 routed expert + 1 个 shared expert，每个 token 只激活 top-k 个 routed，shared 永远参与计算。V4-Flash 的具体取值为 E=256, k=6, shared=1。

激活参数计算（核对值取自 config.json）：

激活 FLOPs 拆解：每次前向激活 7 个 expert（6 routed + 1 shared），等价于 7 × 4096 × 2048 = 58.8M 中间单元；43 层堆叠后总激活参数约 7.6B（experts 部分） + 4–5B（其他）≈ 13B 激活（与官方宣传一致）。其中 shared expert 永远激活，承担了 1.08B / 13B ≈ 8% 的激活开销——共享专家的设计目的就是"用小成本编码全词表通用知识”，让 256 个 routed 专家可以更细粒度地专门化。

与 V3.2 的关键差异：

为什么把 k 从 8 砍到 6？ V4-Flash 是 Flash 定位（8×H100/2×H200 就能跑），参数总预算收缩了 35%（37B→13B 激活）。k 砍 2 个，每层 routed FFN 节省 2/8=25% 的激活 FLOPs，是压缩总参数的最大单点。但 k 太小又会让 routed 容量不足——V4 配套引入 routed_scaling_factor=1.5 来放大 routed 输出（即 “用 6 个专家达到 8 个专家的等效表达力”），这是 V3 没有的"剪枝 + 缩放"组合调参。V3.2-Exp 时代并没有"flash 部署"这个约束，因此维持 k=8。

为什么前 3 层用 hash routing？ V4 的 num_hash_layers=3 是 DeepSeek 在大规模实验后做出的取舍：(1) 浅层 attention 主要在拼词法 / 局部句法，对专家细粒度要求不高，hash routing 分配确定性强、首 token 计算省去一次 gate matmul（4096×256 矩阵-向量乘）；(2) 深层 attention 已经在做长程推理 / 知识调用，必须靠 score routing 选最相关的专家。这 3 层 hash + 40 层 score 的组合在论文 ablation 中略优于"全 score"或"前 6 层 hash"——可视为 V4 的"冷启动"技巧。

routed_scaling_factor=1.5 的角色：这个超参是 V4 独有的关键调参。如果不放大 routed 输出，shared expert 永远在线（≈8% 激活）会"压住" routed 信号；放大 1.5× 等价于让每个 routed expert 的有效权重从 1/6 提升到 1.5/6=0.25，使 routed 加权和 ≈ 1.5（不归一化），shared expert 再叠加 ≈1，总输出幅度合理。这一项不调会出现"routed 加权和 = 0.3 + shared = 1.0，shared 主导"的病态情况。1.5 是 V4 论文 ablation 的最优值。

末尾过渡：以上 256 routed / 6 top-k / 1 shared / hash + score 双路径是 V4-Flash MoE 的"骨架"。§4.2 解释 Gating 网络怎么把 token x 变成 top-6 专家索引的——这是 MoE 的"灵魂"，也是 sqrtsoftplus 与 aux-loss-free 偏置两大 V3-V4 增量的登场之处。

4.2 Gating 网络（MoE 路由）：sqrtsoftplus + Top-k + Aux-Loss-Free Bias

Gating 是 MoE 的"调度中心"：给定一个 token 的 hidden state x ∈ ℝ^{d}（V4-Flash 中 d=4096），Gating 算出 256 个 routed expert 各自的"亲和度分数"，挑出 top-6，并把它们组合回一个加权和。V4-Flash 的 Gating 流程有 5 步，对应源码 inference/model.py:L546-L586 的 Gate 类。

形式化（公式 4.1–4.3）：

(4.1) Gating 评分（不归一化的原始亲和度）：

g_i = sqrtsoftplus( (W_g · x)_i )   ∈ ℝ,    W_g ∈ ℝ^{d × 256}

其中 sqrtsoftplus(z) = sqrt(log(1 + exp(z))) —— 注意是 softplus 后开方，不是 sigmoid、也不是 softmax。config.json 中的 scoring_func: "sqrtsoftplus" 即此。源码 Gate.forward 走 scores = F.softplus(scores).sqrt()（model.py:L571）。

为什么不选 sigmoid/softmax？ V4 论文 ablation 显示：sigmoid 把分数压到 (0,1)，对负 logits 不敏感；softmax 让一个 expert 主导，破坏"细粒度"假设；sqrtsoftplus 在 z → -∞ 时趋近 0，在 z → +∞ 时趋近 sqrt(z)，与 softplus 单调性一致但数值范围更稳——是 routing weight 与 topk 分数同源的最简选择。

(4.2) Top-k 选择（带 aux-loss-free bias）：

score_with_bias = g + b                         ∈ ℝ^{256},  b ∈ ℝ^{256} (trainable, fp32)
topk_idx        = top_k(score_with_bias, k=6)   ∈ {0..255}^6

关键：bias b 仅影响 topk 选择，不参与 routing weight 计算。源码 Gate.forward 把 original_scores = scores（不带 b）单独保存，最后用 original_scores.gather(...) 计算权重（model.py:L580）。

(4.3) Routing weight + 缩放：

weights = g[i for i in topk_idx]    # 取出 top-6 个原始分数（不带 bias）
weights /= weights.sum(...)         # 仅归一化，不 softmax
weights *= self.route_scale         # 乘 1.5

对 sqrtsoftplus/sigmoid 评分，源码会做 weights /= weights.sum(...)（L582），避免分母为 1 时归一化前后无差异；最后 weights *= self.route_scale（L583），route_scale = 1.5。

9 步数据流（每个 token 的生命周期）：

1. 输入 x ∈ ℝ^{4096}，来自上一 Block 的输出（经 mHC 残差混合后）
2. Gate 矩阵乘：W_g · x → ℝ^{256}，W_g 形状 [256, 4096]（= 1M 参数）
3. 评分函数：F.softplus(...).sqrt() —— 整个 256 维过一遍
4. 保存原始分数：original_scores = scores（后续 gather 用）
5. Aux-loss-free bias 加和（仅 topk 路径用）：scores_for_topk = scores + self.bias
6. Top-k 选择：torch.topk(scores_for_topk, k=6, dim=-1)[1] → 6 个 expert 索引
7. Routing weight 提取：weights = original_scores.gather(1, indices) → shape [T, 6]
8. 归一化 + 缩放：weights /= sum 再 *= 1.5
9. 返回 (weights, indices) 给上层 MoE.forward 用

Gate 类的双模式分支：

注：V3 实际用 softmax 评分（scoring_func="softmax"），V3.2-Exp 改用 sigmoid（scoring_func="sigmoid"），V4 用 sqrtsoftplus（scoring_func="sqrtsoftplus"）——这是三阶段连续演化的"评分函数替换"主线。

源码 Gate.__init__ 关键设计（model.py:L546-L565）：

• self.hash = layer_id < args.n_hash_layers —— 一个 bool，前 3 层 True
• 如果 hash 模式（前 3 层）：构造 self.tid2eid ∈ ℝ^{vocab × 6} 的查找表（requires_grad=False），bias 设为 None
• 如果 score 模式（其余 40 层）：构造 self.bias ∈ ℝ^{256}，可训练 fp32 参数

Gate.forward 中两模式分支：

• hash 模式：indices = self.tid2eid[input_ids] —— 一次查表，不算 score
• score 模式：indices = scores.topk(self.topk, dim=-1)[1] —— 算完整 score 后 topk

num_hash_layers=3 这一项是 V4-Flash 相对 V3 的核心工程增量：前 3 层用 hash routing 省去一次 [T, 4096] × [4096, 256] = T × 256 的矩阵乘 + softmax + topk，换来"小批量首 token 速度"提升；同时 hash routing 完全确定（相同 input_id 永远分到相同 expert），让 warmup 阶段负载更稳定。

与 V3.2-Exp 的对比：

V3.2-Exp 用 sigmoid 评分，V4 用 sqrtsoftplus——这是一个不起眼但能"少几次训练震荡"的工程改动。routed_scaling_factor=1.5 则是 V4 配合 k 砍到 6 后的"等效容量补偿"：6 个 routed expert 乘 1.5 等效于 9 个 routed expert（不归一化）的能力。

视觉化：

图 4.1 给出 V4-Flash MoE 完整数据流：(a) token x 经 Gate 矩阵乘 + sqrtsoftplus 评分 → top-6 expert 索引；(b) 256 routed expert 中只有 6 个被激活（深蓝高亮 e4, e48, e100, e156, e205, e231），其余 250 个完全跳过（白色）；(c) 1 shared expert 永远在线（独立虚线箭头）；(d) 输出 = shared + 1.5 × 加权和。底注说明前 3 层走 hash routing 路径。

末尾过渡：以上 sqrtsoftplus 评分 + aux-loss-free bias 加和 + top-6 选专家 + 1.5 缩放构成 Gating 完整流程。但这一切的前提是"score routing"——而 V4 新增了 hash routing 路径，且 aux-loss-free bias 怎么动态更新、训练时和推理时有什么区别，这些是 §4.3 的主题。

4.3 双路径路由 + Aux-Loss-Free 负载均衡

§4.2 把 score routing 路径拆得很细，但 V4-Flash 的 MoE 实际有两条并行的路由路径——前 3 层走 hash、其余 40 层走 score；并且 score 路径的 aux-loss-free bias 有一套独立的更新规则，只在训练时激活。§4.3 把这两件事讲清楚。

4.3.1 双路径路由：按层数分支

V4-Flash 的 43 层 MoE 不是"统一行为"——前 3 层用 hash routing，剩下 40 层用 score routing。这一选择由 config.json 中的 num_hash_layers=3 决定，由源码 Gate.__init__ 中 self.hash = layer_id < args.n_hash_layers 落地（model.py:L554）。

Hash routing 路径（前 3 层）：

• 输入：仅 input_ids（token id），不需要 x 的 hidden state
• 查找表：self.tid2eid ∈ ℝ^{vocab × 6}，形状 [129280, 6]，int32，不可训练（requires_grad=False，model.py:L560）
• 索引：indices = self.tid2eid[input_ids] —— 一次 Python 风格的张量 gather
• 权重：weights = 1.5 / 6 = 0.25（均匀权重，无 softmax 归一化）
• 计算代价：0 个 matmul、0 个 softmax、0 个 topk

为什么用 hash routing？ 两个原因。(1) 性能：前 3 层如果用 score routing，每次前向都要算 x @ W_g（4096×256 矩阵乘 + sqrtsoftplus + topk），对 batch=1 的 decode 阶段这是非平凡延迟；hash 一次查表就完成，延迟降一个数量级。(2) 稳定性：浅层 attention 主要在拼词法/局部句法，对"哪个专家擅长哪个子任务"的要求不高；hash 强制把 token id 平均分布到 256 个 expert（前提是哈希函数设计良好），让浅层 expert 不至于被某些高频 token 主导。论文 ablation 显示：前 3 层用 hash + 后续 40 层用 score 的组合在 PPL/MMLU 上略优于"全 score"或"前 6 层 hash"。

Score routing 路径（后 40 层）：

• 走 §4.2 的完整 5 步：W_g · x → sqrtsoftplus → + b → top-6 → softmax 权重
• 训练时：bias b 每步更新（见 §4.3.2）
• 推理时：bias 冻结，与 hash 路径行为一致（用训练好的 b 算 topk）

两路径在源码中的统一：Gate.forward 末尾只返回 (weights, indices)，hash 和 score 路径在调用方 MoE.forward 看来完全一致——expert 调度与组合逻辑共享。这也是为什么 V4-Flash 能在 43 层中无缝混用：MoE 层只看到"6 个 expert 索引 + 6 个权重"，不关心是 hash 来的还是 score 来的。

num_hash_layers=3 的设计依据：V4 选 3 是"浅层 token 嵌入靠哈希足够、深层 expert 路由必须靠 score"的经验折中。0 层 hash（全 score routing）会拖慢前 3 层 batch=1 decode 延迟一个数量级；> 6 层 hash 会显著损害 PPL，因为深层 attention 高度依赖 score 选择最相关 expert。V4 论文未公开具体 ablation 数字，定性结论是"3 是甜点，速度提升约 10% 而 PPL 退化可忽略"。

4.3.2 Aux-Loss-Free 负载均衡

V4-Flash 没有传统 aux loss（λ · Σ_i f_i · p_i 那种），而是用 aux-loss-free 偏置（沿用 V3.2-Exp 论文 §3.2 设计）。这一设计的核心动机是：aux loss 会污染主损失梯度，让模型"为了均匀而均匀"牺牲性能；aux-loss-free 则用"训练时动一动手、推理时不影响"的偏置 trick，避开这个矛盾。

核心规则（每 step，源码 MoE.forward 训练分支未在本快照中——本仓库 inference-only，训练时偏置更新在训练框架侧）：

1. 每个 routed expert e 维护一个可训练偏置 b_e ∈ ℝ（fp32，形状 [256]）
2. 训练 step 结束时，统计该 step 中各 expert 被 top-6 选中的频率 f_e
3. 目标频率 p = 1/256 ≈ 0.0039
4. 偏差更新（论文公式 §4.2.3 L1611–L1649）：

   b_e ← b_e - η · (1/E - f_e)         (E=256, η=0.001, 论文指定)
   

   等价描述：
   • 若 f_e > p + δ（过载）：b_e -= η → 下一 step topk 选该 expert 的概率下降
   • 若 f_e < p - δ（欠载）：b_e += η → 下一 step topk 选该 expert 的概率上升
   • δ 是"容忍带"，论文给 δ≈0.05
5. 关键约束：b_e 仅影响 topk 选择（公式 4.2 中的 score_with_bias），不参与 routing weight（公式 4.3 中的 original_scores.gather(...)）。源码 Gate.forward 保留 original_scores = scores（L572），最后用 original_scores.gather(1, indices)（L580）取权重——bias 完全不进 softmax 分母。

为什么"aux-loss-free"比"aux loss"好？

V3 论文 §3.2 报告：在 200B token 训练中，aux-loss-free 与 aux loss 达到几乎相同的负载均衡（最大 expert 频率 0.012 vs 0.011），但下游任务平均提升 +0.5% MMLU——负载均衡不一定要从主损失里"偷梯度"。

源码对应：本仓库是 inference-only，不包含训练时偏置更新的代码。训练时偏置更新由 DeepSeek 训练框架（不在本快照内）实现，按 V3 论文算法：每 step 末尾统计 bincount(indices) / total_tokens，按上述规则做 b_e -= η · (1/256 - f_e)。本快照内可看到 Gate.__init__ L562 定义偏置参数 + Gate.forward L574-L575 加和（仅 score 模式）。

4.3.3 视觉化

图 4.2 给出 Gating 完整双路径流程。左侧蓝色路径是 score routing（40 层）：W_g · x → sqrtsoftplus → + b → top-6 → softmax 权重；右侧橙色虚线框是 hash routing（前 3 层）：hash(input_id) → tid2eid[input_id] → 均匀权重 1.5/6。中间底部是 aux-loss-free bias 的训练时更新循环：[a] 统计 → [b] 与 p=1/256 比较 → [c] 按 η=0.001 步长更新 b_e。两条路径在 MoE.forward 入口汇合，输出 y = shared + 1.5 × 加权和。

末尾过渡：以上双路径路由 + aux-loss-free 偏置更新是 V4-Flash MoE 的"运行机制"。但训练时和推理时的行为还有几处细节差异（FP4 反量化时序、MTP 是否参与、expert 调度策略），§4.4 展开这些运行时差异，并给出负载均衡的实测效果。

4.4 训练 vs 推理差异 + 负载均衡实测

V4-Flash MoE 的"前向计算"在训练和推理时大体一致——同一份 Gate/Expert/MoE 类、同一套 sqrtsoftplus 评分、同一个 top-6 选择。但有几个关键差异值得展开，因为它们直接影响部署策略与训练调参。

4.4.1 训练时 vs 推理时的 MoE 行为对照