[{"content":"","date":"2026年8月4日","externalUrl":null,"permalink":"/tags/grpo/","section":"Tags","summary":"","title":"GRPO","type":"tags"},{"content":"","date":"2026年8月4日","externalUrl":null,"permalink":"/tags/llm/","section":"Tags","summary":"","title":"LLM","type":"tags"},{"content":"","date":"2026年8月4日","externalUrl":null,"permalink":"/tags/ppo/","section":"Tags","summary":"","title":"PPO","type":"tags"},{"content":" 原论文：Reinforcement Learning Meets Large Language Models: A Survey of Advancements and Applications Across the LLM Lifecycle，Keliang Liu 等，arXiv:2509.16679（2025-09）。\n阅读地图 # 这篇综述把 RL 贯穿 LLM 全生命周期，主线是三分法：\n预训练（Pre-training / Mid-training）：为后续 RL 做数据与风格适配，或把 next-token prediction 改造成可验证奖励的 RL 任务 对齐微调（Alignment）：RLHF、偏好优化、奖励模型设计，对齐人类意图与安全 强化推理（RLVR）：用可自动验证的客观奖励（数学判分、单测等）推高推理上限——综述强调这是近年主战场 相对只谈 RLHF 的旧综述，本文额外覆盖了数据集与基准、开源工具框架，以及 RLVR 能力边界、熵坍缩等争议话题。\nPreliminaries：Policy 与 Value 的地基 # MDP 与两大范式 # RL 通常建模为 MDP：状态 $s$、动作 $a$、转移、奖励 $r$。目标是学到策略 $\\pi$，最大化期望累积回报。\nPolicy-based：直接优化 $\\pi_\\theta(a|s)$（策略梯度族） Value-based：估计 $V$ 或 $Q$，再由值函数导出策略（Q-learning / SARSA / DQN） LLM 场景的动作空间约等于词表（乃至整段序列），很难为每一个可能输出维护显式的 $Q$ 值——这正是 value-based 方法很少作为 LLM RLHF 主框架的原因：像表格型或小离散动作环境那样为每个动作维护 Q 值，在词表级动作空间下不现实。因此 RLHF / RLVR 的主流是 policy gradient（PPO、GRPO 等），而非 DQN 主框架；value 思想仍会以 critic 或 baseline 的形式出现。\nPolicy Gradient、Baseline、优势 # 目标是 $J(\\theta)=\\mathbb{E}_{\\tau\\sim\\pi_\\theta}[R(\\tau)]$，写成期望形式：\n$$ J(\\theta)=\\sum_\\tau \\pi_\\theta(\\tau)R(\\tau) \\quad\\Rightarrow\\quad \\nabla J=\\sum_\\tau \\nabla\\pi_\\theta(\\tau)R(\\tau) $$直接对 $\\pi_\\theta(\\tau)$ 求导不好蒙特卡洛估计，这里用 log-derivative trick 这一关键恒等式：\n$$ \\nabla_\\theta\\pi_\\theta(\\tau)=\\pi_\\theta(\\tau)\\nabla_\\theta\\log\\pi_\\theta(\\tau) $$代入后\n$$ \\nabla J=\\mathbb{E}_{\\tau\\sim\\pi_\\theta}\\big[\\nabla\\log\\pi_\\theta(\\tau)R(\\tau)\\big] $$轨迹概率取对数后乘积变求和，就得到综述公式 (1)：\n$$ \\nabla_\\theta J(\\theta)=\\mathbb{E}_{\\tau\\sim\\pi_\\theta}\\Big[\\sum_t\\nabla_\\theta\\log\\pi_\\theta(a_t|s_t)R_t\\Big] $$这里 $\\log$ 不改变优化目标，而是把「对概率求导」变成「对已采样动作的 log 概率求导」，从而可以用采样轨迹做蒙特卡洛估计。\n引入一个仅依赖状态的 baseline $b(s)$（公式 (2)）：\n$$ \\nabla_\\theta J(\\theta)=\\mathbb{E}\\Big[\\sum_t\\nabla_\\theta\\log\\pi_\\theta(a_t|s_t)\\,(R_t-b(s_t))\\Big] $$$A_t=R_t-b(s_t)$ 即优势。为什么减去 $b(s)$ 不引入偏差？对固定 $s$，$b(s)$ 与动作无关：\n$$ \\sum_a \\pi(a|s)\\nabla\\log\\pi(a|s)b(s)=b(s)\\nabla\\Big(\\sum_a\\pi(a|s)\\Big)=b(s)\\nabla 1=0 $$即 $\\mathbb{E}_{a\\sim\\pi}[\\nabla\\log\\pi\\cdot b(s)]=0$，减去 baseline 不改变梯度期望；若 $b(s)$ 与回报正相关，则能显著降低方差——这是无偏但可降方差的直觉来源。\nActor-Critic 结构里，actor 更新策略，critic 估计 $V$ 或 $Q$，以提供低方差的优势估计（或 TD 误差）。\nTRPO → PPO # TRPO 在 KL 信任域约束下最大化优势期望（公式 (3)）。PPO 用 clipped surrogate 做工程近似（公式 (4)）：\n$$ L^{\\mathrm{PPO}}(\\theta)=\\mathbb{E}_t\\Big[\\min\\big(r_t(\\theta)\\hat{A}_t,\\;\\mathrm{clip}(r_t(\\theta),1-\\epsilon,1+\\epsilon)\\hat{A}_t\\big)\\Big] $$其中 $r_t(\\theta)=\\pi_\\theta/\\pi_{\\mathrm{old}}$。目标外层取 $\\min$ 是一个悲观下界，保证不会因为「涨分」而鼓励过大的更新：\n$\\hat{A}_t\u003e0$：想提高这个动作的概率；当 $r\u003e1+\\epsilon$ 后 clip 封顶，$\\min$ 阻止好动作概率涨过头 $\\hat{A}_t\u003c0$：想降低这个动作的概率；当 $r\u003c1-\\epsilon$ 后 clip 封底，$\\min$ 阻止差动作概率砍过头 GRPO：去掉 Critic，用组内相对奖励 # LLM 推理场景下 PPO 有几个痛点：(1) 额外的 value 网络带来显存和算力开销；(2) 长序列上 value 估计不准；(3) 传统设置常一次只打一条回复，学习效率低。\nGRPO（DeepSeekMath 提出）的做法是：每道题采样 $G$ 条回复成组，用奖励模型或规则打分；去掉独立的 critic，用组内平均奖励作动态 baseline：$A_i=R_i-\\bar{R}_{\\mathrm{group}}$。\n综述正文明确写的是 group average，但印刷公式 (5) 排版成了 $\\max$，这与正文叙述以及 DeepSeekMath 原式都不一致——应按 mean / std 理解：\n$$ \\hat{A}_{i,t}=\\frac{r_i-\\mathrm{mean}(\\{R_i\\}_{i=1}^{G})}{\\mathrm{std}(\\{R_i\\}_{i=1}^{G})} $$分子是该条奖励相对组统计量的偏离，分母是组内标准差（做尺度归一，近似 z-score）。若真的按 max 理解，多数样本的优势会 $\\le 0$，这与「组内相对」的设计初衷不符，因此印 max 几乎可以肯定是笔误。全组奖励相同时 std 会趋近于 0，此时应让优势为 0（没有相对信号可言），实现上通常加一个小 $\\varepsilon$ 防止除零。\n公式 (6) 在 clip 目标上还做了两处调整：对 token 长度 $|o_i|$ 做平均，以及直接在目标中加上 $-\\beta D_{\\mathrm{KL}}(\\pi_\\theta\\|\\pi_{\\mathrm{ref}})$（而经典 PPO-RLHF 更常把 KL 折进 reward 里）。这两种放置方式的差异值得展开：把 KL 折进 reward 会改变有效回报，进而影响优势与 baseline 的关系，容易「污染」优势的尺度；GRPO 把 KL 作为独立正则项加进 loss，优势仍完全由任务奖励（组内相对）计算，更「干净」，也更容易和组内相对的设计保持自洽。\nPPO vs GRPO 对照\n维度 PPO（经典 RLHF） GRPO 采样 常一次一条 每题 $G$ 条成组 优势 value/critic（+GAE） 组内相对奖励，无独立 critic KL 常进 reward 常直接进损失 目标骨架 $\\min(r\\hat{A},\\mathrm{clip}\\cdot\\hat{A})$ 同结构 + 组平均 + 长度归一 除以 std 这个设计本身也有取舍：好处是尺度无关、能对难度自适应；风险是 std 很小时信号会被放大到不稳定，稀疏 0/1 可验证奖励下组内方差结构比较特殊，整组全对或全错时更是完全没有相对信号。常见的改法包括 std 过小时把优势直接置零、只减均值不除 std（部分 R1-Zero / Dr.GRPO 一脉的讨论）、过滤掉零方差的组，或者混入过程奖励增加区分度。\nValue Learning 速记 # Q-learning（公式 (7)）：用 $\\max_{a'}Q(s',a')$，行为策略可以不同于目标策略 → off-policy SARSA（公式 (8)）：用实际执行的 $a_{t+1}$ → on-policy DQN（公式 (9)）：神经网络逼近 $Q$，配合目标网络与 replay buffer 预训练与对齐 # 预训练阶段的 RL 与 Mid-training # 多数 RL 工作仍落在对齐与后训练，但预训练侧也有一些探索：\nReinforcement Pre-Training：把 next-token 预测改造成「正确预测下一个 token 得可验证奖励」的 RL 推理任务，资源较重，常需要已具备推理能力的底座 视觉预训练：把无标注图像预训练也框成 RL 问题（如 Annotation Bootstrapping） OctoThinker 的 Mid-training：两阶段 mid-training 提升底座与后续 RL 的兼容性，例如让原本不太适合 RL 的 Llama 在数学推理上追平同量级的 Qwen Mid-training（综述 §3.1）在训练方式上仍是 next-token / next-word 自监督，和预训练没有区别；不同的是目标——把预训练模型改造得更适合后续 RL，数据也从海量杂文转向高质量、任务相关的语料。数据质量、风格与课表调度，对 RL 阶段能否 scaling 起来很关键。\n经典对齐：RLHF 与偏好优化 # InstructGPT 范式常分三阶段：\nSFT：在演示数据上做监督微调，得到可遵从指令的初始策略 Reward Modeling：在偏好对上训练打分模型 $r_\\phi(q,o)$ RL（常用 PPO）：用 $r_\\phi$ 打分，加上相对 reference 的 KL 约束，优化策略 相关工作还涵盖信息论视角的迭代对齐、Constitutional AI / RLAIF（用 AI 反馈替代人类标注），以及缓解 reward hacking 的方法——策略学会利用奖励函数的漏洞刷高代理奖励，但真实的有用性或正确性并未提升，在 RM 不完美或规则存在漏洞时尤其常见，是对齐与 RLVR 共同的核心风险。\nDPO 在一定假设下绕过了显式的奖励建模与 RL 优化，直接用偏好数据把策略微调到偏好诱导的最优策略；传统 RLHF 才是「先训 RM 再跑 PPO」的两段式。后续还有 KTO、ORPO、$\\beta$-DPO、ΨPO 等变体。\n奖励模型的新方向 # 推理式 / 生成式 RM：在打分时引入测试时算力、CoT，甚至代码验证（RRM、GenPRM、RM-R1 等） 过程奖励（PRM）：对推理中间步骤打分，而不是只给终局一个稀疏结果奖励——信号更密集，但标注和「步骤对错」的定义都更难 原则 / 规则奖励：用自然语言原则、或从偏好数据中抽取规则，再经 verifier 度量满足程度（RewardAnything、AUTORULE 等） 训练环中，奖励模型（或规则）只负责产出标量或过程奖励 $r$，本身并不直接对策略做监督梯度更新——真正驱动策略更新的是优势与 PPO/GRPO 目标：\nprompt → π 采样回复(s) → RM/规则打分得 r → 算 Â（PPO: critic；GRPO: 组内相对）→ clip 目标更新 π DPO 路径则可以绕开这个显式 RM + RL 的训练环。\nRLVR：强化推理 # 定义与能力边界之争 # RLVR（Reinforcement Learning with Verifiable Rewards）用程序检查、数学答案验证、形式化证明等可自动验证的客观奖励做 RL 微调，不必依赖人类偏好训练出的奖励模型（也可以和规则奖励并存）。这是 o1 / R1 一路方法的核心范式之一。\n综述里一个尚无定论的争议是：RL 究竟是真正扩展了推理能力，还是只是放大了底座模型分布里本就存在的高奖励路径？\nYue et al. 用 pass@k（采样 $k$ 次、至少一次正确的概率）做评估，发现小 $k$（如 $k=1$）时 RLVR 更好，但 $k$ 增大后 base 模型往往反而更好；结合覆盖度和困惑度分析，他们认为 RLVR 的正确轨迹大多落在 base 分布内，更像是提高了采样效率，能力边界甚至有所收窄，模型的反思行为也常能在 base 模型里找到源头。 Liu / ProRL 等的工作则显示，在足够的训练时间和新任务上，RL 能发现 base 模型完全没有的新解法路径，据此主张 RL 确实可以扩展能力边界。 Wu 等认为 RLVR 主要还是一个高效采样器，偶尔会越界，但训练中多样性坍缩和遗忘现象同样会发生。 综述本身并列了双方证据，没有给出单方面的定论。一个可能的折中判断是：区分「pass@1 的实用提升」和「能力边界 / 新颖推理路径」这两层不同的主张——前者证据相当扎实，后者仍存在争议，取决于训练时长、任务新颖度等具体条件。\n熵、性能上限与训练动态 # Cui et al. 给出了熵 $H$ 与下游性能 $R$ 之间的一个经验关系（公式 (10)）：\n$$ R=-a\\exp(H)+b $$$H$ 下降时，$\\exp(H)$ 随之下降，$-a\\exp(H)$ 上升，$R$ 也随之上升（在拟合成立的区间内）；当 $H\\to 0$ 时，$R\\to -a+b$，给出了一个可预测的性能上限。这意味着性能提升常常以消耗熵为代价，靠不断压低熵来换性能不可能无限持续——熵一旦耗尽，就会限制「只靠堆算力做 RL」的收益，需要熵管理，或者更好的探索 / 目标函数设计，而不只是简单加一个熵正则项。\n**熵坍缩（entropy collapse）**指策略分布越来越尖锐、生成多样性下降的现象。它和「遗忘预训练知识」相关但不完全等同——综述在这里强调的重点是探索耗尽与扩展上限的关系。高协方差、高熵的 token（例如逻辑连接词）对整体熵和推理路径的影响尤其大，限制更新对象、设计优于朴素熵损失的目标函数，都是目前的活跃方向。\n算法进展：DAPO 等 # DAPO（基于 GRPO）针对长 CoT 场景做了几处改动：Clip-Higher、Dynamic Sampling、Token-Level Policy Gradient Loss、Overlong Reward Shaping。\n标准的对称 clip 会限制低概率动作的更新幅度，容易促成熵坍缩；Clip-Higher 的做法是放宽 clip 的上界，以保留更多探索空间。此外还包括对极端奖励的 prompt 做动态采样、按序列长度加权、对过长输出做惩罚或截断，用来缓解样本浪费和输出冗长失控的问题。\n另外还有树结构 RL、对抗 / 多智能体训练、序列级重要性采样（GSPO）等分支，这里从略。\n多模态、自适应长度与 Agent # 多模态：视觉、视频、具身等场景常有显式 ground truth，比较适合规则奖励；难点在于可验证性弱于数学、容易出现文本偏见而忽略图像信息，以及过程稠密奖励的设计更复杂 自适应推理：控制 CoT 长度以及是否需要思考，避免简单题过度思考、难题预算不足 Agent：多轮交互下延迟或稀疏的终局奖励（如 Echo Trap 现象）、工具调用带来的动作空间、长程信用分配、记忆管理；过程奖励和结果奖励哪个更合适也更难定义（如 SPA-RL、LARM） 把 RLVR 从数学 / 代码推广到多模态推理或多轮 Agent 时，奖励设计会遇到几类新困难：可验证性上，视觉、空间、具身类的结果很难像数学答案那样规则化判分；稀疏和延迟上，多轮工具调用、长程任务的终局奖励天然稀疏且滞后；过程与结果的取舍上，中间的工具调用步骤对错本身就难标注，因此需要更细粒度的密集奖励、过程奖励或进度分解设计（例如宏动作层面的优势估计、检索 token 的 mask、专门的记忆管理机制）。\n数据、工具与开放问题 # 数据与基准类型 # 对齐 / 对话偏好：人类标注偏好、AI 反馈偏好 可验证语料：数学（GSM8K、MATH、AIME 等）、代码（LiveCodeBench、SWE-bench 等）、程序 / 形式化验证式奖励 合成数据：任务定义驱动的合成 RL 数据、多步工具合成等 评测：对齐基准 + 推理基准（含 Agent / 多轮工具） 工具与框架 # 综述整理了主流开源 RLHF / RLVR 训练框架与库（trlX、各厂商开源训练栈等），实践中按底座模型与分布式 rollout 的需求选型即可；具体细节随生态更新较快，以综述原文和各框架官方文档为准。\n开放问题 # 综述列出的开放问题包括：能力边界是否真正扩展、熵管理与探索、奖励黑客、过程奖励质量、Agent 长程信用分配、RLIF（内部反馈）后期退化、多模态可验证奖励设计等，这些都可以进一步拆解成具体、可检验的研究假说。\n公式速查 # 公式 内容 (1) $\\nabla\\log\\pi\\cdot R_t$（REINFORCE） (2) baseline / advantage：$R_t-b(s_t)$ (4) PPO $\\min(r\\hat{A},\\mathrm{clip}\\cdot\\hat{A})$ (5) 印刷为 $r_i-\\max$ over std；正文与原论文应理解为 mean (6) GRPO clip + 组相对优势 + KL 项 + 长度平均 (10) $R=-a\\exp(H)+b$（熵–性能） ","date":"2026年8月4日","externalUrl":null,"permalink":"/notes/papers/rl-meets-llms-survey/","section":"笔记","summary":"梳理 RL 如何贯穿 LLM 全生命周期——预训练/Mid-training、RLHF 对齐微调、RLVR 强化推理三条主线；从策略梯度、PPO 到 GRPO 的算法演进，以及 RLVR 是否真正扩展推理能力、熵坍缩与性能上限等争议。","title":"RL Meets LLMs：大语言模型全生命周期强化学习综述","type":"notes"},{"content":"","date":"2026年8月4日","externalUrl":null,"permalink":"/tags/rlhf/","section":"Tags","summary":"","title":"RLHF","type":"tags"},{"content":"","date":"2026年8月4日","externalUrl":null,"permalink":"/tags/rlvr/","section":"Tags","summary":"","title":"RLVR","type":"tags"},{"content":"","date":"2026年8月4日","externalUrl":null,"permalink":"/tags/","section":"Tags","summary":"","title":"Tags","type":"tags"},{"content":"","date":"2026年8月4日","externalUrl":null,"permalink":"/tags/%E5%BC%BA%E5%8C%96%E5%AD%A6%E4%B9%A0/","section":"Tags","summary":"","title":"强化学习","type":"tags"},{"content":"CUDA 与 GPU 系统相关笔记。\n","date":"2026年7月10日","externalUrl":null,"permalink":"/notes/cuda/","section":"笔记","summary":"","title":"CUDA","type":"notes"},{"content":"","date":"2026年7月10日","externalUrl":null,"permalink":"/tags/deepseek-v4/","section":"Tags","summary":"","title":"DeepSeek V4","type":"tags"},{"content":"","date":"2026年7月10日","externalUrl":null,"permalink":"/tags/kv-cache/","section":"Tags","summary":"","title":"KV Cache","type":"tags"},{"content":"LLM 推理优化相关笔记。\n","date":"2026年7月10日","externalUrl":null,"permalink":"/notes/llm-inference/","section":"笔记","summary":"","title":"LLM 推理","type":"notes"},{"content":"","date":"2026年7月10日","externalUrl":null,"permalink":"/tags/llm-%E6%8E%A8%E7%90%86/","section":"Tags","summary":"","title":"LLM 推理","type":"tags"},{"content":"","date":"2026年7月10日","externalUrl":null,"permalink":"/tags/mla/","section":"Tags","summary":"","title":"MLA","type":"tags"},{"content":"","date":"2026年7月10日","externalUrl":null,"permalink":"/tags/pagedattention/","section":"Tags","summary":"","title":"PagedAttention","type":"tags"},{"content":"Triton kernel 编程与优化笔记。\n","date":"2026年7月10日","externalUrl":null,"permalink":"/notes/triton/","section":"笔记","summary":"","title":"Triton","type":"notes"},{"content":" ","date":"2026年7月10日","externalUrl":null,"permalink":"/notes/","section":"笔记","summary":"","title":"笔记","type":"notes"},{"content":"分布式训练与推理系统笔记。\n","date":"2026年7月10日","externalUrl":null,"permalink":"/notes/systems/","section":"笔记","summary":"","title":"分布式系统","type":"notes"},{"content":"我是 Yiwen Cai，北京邮电大学计算机技术专业硕士在读。我的研究兴趣集中在算法与硬件之间的那一层——GPU 算子优化与 LLM 训推优化。目前在腾讯混元 AI Infra 实习，做推理框架调度与高性能算子优化。\n教育经历 # 北京邮电大学 · 硕士（2026.09 – 2029.06）计算机技术（推免） 实验室：ParCIS 并行计算与智能系统实验室 导师：李士刚 教授 北京邮电大学 · 本科（2022.09 – 2026.06）计算机科学与技术 GPA 3.76 / 4.0 · Rank 59/389 · CET6 552 实习经历 # 腾讯混元 AI Infra（2026.07 – 至今）—— LLM 推理框架调度与高性能算子优化 覆盖 Blackwell 架构 attention / 量化 / 融合类算子，主要开源仓库 Tencent/hpc-ops 部分贡献已合入 vLLM main 分支 研究方向 # LLM 训推优化 GPU 算子优化（CUDA / Triton） 分布式计算与云计算、高性能计算 代表性工作 # 华为 ICT 大赛全国总决赛一等奖（最高算子迁移优化奖）——64 进 8、决赛第 3；作为队伍唯一开发者获算子迁移优化赛题所有队伍最佳。 cuda-gemm——从零手写 CUDA GEMM；SGEMM 达 cuBLAS 95.1%，HGEMM Tensor Core 达 213 TFLOPS。 FastGNN——基于 Tensor Core 的 GNN 训练系统；相对 DGL/PyG 加速 1.8x–3.3x。 第十五届蓝桥杯北京赛区二等奖。 详见 项目 页。 简历 # 中文简历（PDF） English Resume (PDF) 联系 # GitHub：@yiwen-cai 邮箱：caiyiwen.cs@foxmail.com ","date":"2026年7月10日","externalUrl":null,"permalink":"/about/","section":"首页","summary":"","title":"关于","type":"page"},{"content":"LLM 推理加速、KV Cache、注意力机制等方向的论文阅读笔记。\n","date":"2026年7月10日","externalUrl":null,"permalink":"/notes/papers/","section":"笔记","summary":"","title":"论文阅读","type":"notes"},{"content":" ","date":"2026年7月10日","externalUrl":null,"permalink":"/projects/","section":"项目","summary":"","title":"项目","type":"projects"},{"content":"","date":"2026年7月10日","externalUrl":null,"permalink":"/tags/%E9%95%BF%E4%B8%8A%E4%B8%8B%E6%96%87/","section":"Tags","summary":"","title":"长上下文","type":"tags"},{"content":" 本文基于 2026 年组会综述汇报整理，覆盖 2023–2026 年 KV Cache 优化的系统管理、缓存压缩与架构协同三条主线，并梳理最新工作（DeepSeek V4 CSA/HCA、MLRA、GDN 等）。\n长上下文与长输出正在成为 LLM 推理的常态——长文档问答、RAG、Agent 工具调用、多轮对话、Reasoning Model 的长链推理，都在把 KV Cache 从「推理加速机制」推向「核心资源瓶颈」。本文不堆砌论文，而是用一条公式作为思维骨架，把纷繁的方法组织成一个有层次的优化图谱。\n核心公式 # $$\\text{KV Cache Size} = 2 \\times L \\times T \\times H_{kv} \\times D_h \\times \\text{bytes}$$ 符号 含义 $2$ key + value 两份 $L$ 层数 $T$ 上下文长度（含已生成 token） $H_{kv}$ KV head 数 $D_h$ 每个 head 的维度 $\\text{bytes}$ 每个元素的字节数 全文每个优化方法，都锚定在这个公式的某个变量上：\n系统管理优化 $T$ 的调度与放置（不减少 $T$ 本身） 压缩 / 量化减少有效 $T$ 或 $\\text{bytes}$ 架构协同减少 $H_{kv}$、改变缓存表示，或压缩有效 $T$ 一、为什么 KV Cache 成为新瓶颈？ # 本章建立动机：在核心公式里，$T$ 正在不断增长。\n自回归推理的两个阶段：Prefill 与 Decode # LLM 推理分为 prefill 和 decode 两个阶段，它们的资源瓶颈截然不同：\nPrefill：把完整 prompt 一次性送入模型，并行处理，构建初始 KV Cache。这一阶段算力密集，更偏 compute-bound。 Decode：每次只生成一个 token，每一步都要读取全部历史 KV Cache，同时追加新 token 的 K/V。访存密集，更偏 memory-bandwidth-bound。 一句话：Prefill builds the cache; Decode consumes and extends the cache.\nPrompt tokens → Prefill（并行处理 / 构建 KV Cache / compute-bound） → Decode step 1（读 cache → 生成 token 1 → 追加 KV） → Decode step 2（读更新后的 cache → 生成 token 2 → 追加 KV） → ...（memory-bandwidth-bound） 长上下文与长输出下，decode 阶段会不断访问越来越大的 cache，显存带宽压力非常突出。\nKV Cache 缓存什么：动态推理状态，不是模型权重 # Transformer 每一层都会为 token 计算 key 和 value（$K = XW_K$，$V = XW_V$）。自回归生成时，历史 token 的 K/V 对后续 step 不会变化，所以可以缓存起来。KV Cache 和模型权重是两类完全不同的东西：\n项目 模型权重 KV Cache 性质 固定参数 动态状态 是否请求相关 否 是 是否可共享 多请求共享 通常每请求独立 增长方式 固定大小 随上下文和输出增长 常见优化 权重量化、剪枝 管理、压缩、量化、架构协同 KV Cache 不是模型权重，而是每个请求在推理过程中动态增长的状态。 这个区分决定了后续优化路线的根本差异。\n显存公式与线性增长 # 回到核心公式 $\\text{KV Cache Size} = 2 \\times L \\times T \\times H_{kv} \\times D_h \\times \\text{bytes}$。最重要的变量是 $T$——cache 随 $T$ 线性增长。serving 场景还要乘上并发请求数，瓶颈从两个方向逼近：\n容量压力：cache 能不能存下（GPU HBM 容量有限） 带宽压力：decode 每一步读取历史 KV 的 HBM 带宽 上下文从 4K → 32K → 128K → 1M，KV Cache 同步线性膨胀，容量和带宽双双吃紧。\n2025–2026 的新压力：三种 Long # 传统长上下文主要关注 long input（长文档、RAG）。但 2025–2026 的 workload 带来了三种新的「Long」：\nLong Input：长文档问答、RAG、多文档检索——压力来自 prefill 后形成的大规模初始 KV Cache。 Long Output：Chain-of-thought reasoning、self-reflection、verification——压力来自 decode 中不断追加新 token 的 K/V。输入可能不长，但 CoT 很长。 Long-lived State：multi-turn dialogue、Agent tool calls——cache 生命周期更长，涉及恢复、迁移、复用。 三者汇聚：更大、更长生命周期、workload-dependent 的 KV Cache → GPU 显存容量 + 带宽瓶颈。\nKV Cache 优化正在从「长输入管理」扩展到「长输出推理」和「长生命周期状态管理」。 这是后文第五部分前沿趋势的伏笔，也是 2025–2026 KV Cache 研究升温的重要原因。\n二、系统管理：KV Cache 怎么放、怎么调度、怎么恢复 # 本章的发力点是 $T$ 的管理与调度，但不减少 $T$ 本身。明确边界：系统管理提高的是利用率，不改变每个 token 的 KV 表示。\nPagedAttention：像虚拟内存一样管理 KV Cache # PagedAttention（vLLM, SOSP 2023）借鉴操作系统的分页思想：请求在逻辑上看到连续的 token 序列，但 KV blocks 在物理 GPU memory 中可以不连续，由一张 block table 做映射。\n逻辑 KV blocks（请求视角：连续） block table（页表） GPU 物理块（分散） block 0 ─────────────────────────→ 0 → 物理 7 ────→ block 7 block 1 ─────────────────────────→ 1 → 物理 2 ────→ block 2 block 2 ─────────────────────────→ 2 → 物理 9 ────→ block 9 block 3 ─────────────────────────→ 3 → 物理 4 ────→ block 4 它解决的是：连续显存分配的碎片、输出长度未知导致的预分配浪费、decode 动态增长，从而支撑更大的有效 batch size 和 continuous batching。\nPagedAttention is memory management, not KV compression. 它不改变 K/V 表示本身。\nKV-aware scheduling：把 cache 容量变成一等调度约束 # PagedAttention 解决了 KV blocks 在 GPU 里怎么组织，但真实服务中多个请求同时到达、长度不同、输出预算不同。传统按到达时间 batching 只关注 batch size 和 compute，忽略未来 KV Cache 占用，可能导致 OOM 或高延迟。\nOnline Scheduling for LLM Inference with KV Cache Constraints（2025）一类的工作把 KV Cache capacity 纳入调度模型：估算 prompt 长度 + 输出预算 + KV footprint，在 KV cache 约束下做在线 batching，换取更低延迟和更好资源利用。\nKV Cache 不只是内存对象，而是 online serving scheduler 必须显式建模的一等约束。\nAgent 与多轮对话：cache 生命周期变长了 # Agent workload 和普通单轮问答不同。Agent 会多次规划、调用工具、等待返回，然后继续推理。这个过程中 KV Cache 的状态会在多种角色间流转：\nactive（正在推理）—— tool call → idle（等待返回） idle —— resume → active；或空间竞争 → evicted evicted —— 重算/加载 → restored active —— prefix 共享 → shared / reused critical agent cache 被挤出会显著增加 TTFT（首 token 延迟）——这是 idle 期间不能随便丢 cache 的根本原因。\n新挑战包括：cache 在工具调用期间 idle、critical agent cache 被挤出会显著增加 TTFT、多 agent 并发导致空间竞争、shared prefix / session cache 需要复用策略。代表方向如 TokenCake（KV-Cache-Centric Serving for Multi-Agent）、Continuum（带 KV Cache TTL 的 multi-turn agent 调度）。\nAgent 场景使 KV Cache 管理从「单次请求资源分配」变成「跨时间的状态生命周期管理」。\nKV Cache restoration：重算还是加载？ # 多轮对话、RAG、Agent 场景经常需要恢复已有 KV Cache。恢复有两种方式，各有代价：\n重算（Recompute）：从原始 prompt 重新计算，消耗 GPU compute，长上下文下代价高。 加载（Load）：从 CPU/SSD/remote 加载，消耗 I/O 带宽，可能增加 TTFT。 CacheFlow 把这个问题变成 token / layer / multi-GPU 三维并行恢复：token 级（早期 chunk 重算、后期 chunk 加载）、layer 级（低层重算、高层加载）、multi-GPU 级（并发恢复 shard）。Kareto 则从 multi-objective tiered storage 配置角度，在 latency ↔ throughput ↔ cost 之间权衡，跨 GPU HBM / DRAM / SSD / remote 做分层。\n2026 年的系统问题不只是「cache 放在哪里」，还包括「cache 如何在计算和 I/O 之间高效恢复」。\n三个常被混淆的层次：FlashAttention / PagedAttention / KV Compression # 这是一个常被混淆的关键区分，三者位于不同层次，不是替代关系：\n层次 代表方法 解决的问题 Attention kernel FlashAttention 减少 attention 计算中的 HBM/SRAM IO，避免 materialize attention matrix KV block management PagedAttention 减少碎片，支持动态 batching Cache policy eviction / quantization / retrieval 减少保存或访问的 KV FlashAttention 优化 attention 计算路径（尤其 prefill/训练的 IO）；PagedAttention 管理 decode serving 中的 KV blocks；KV compression 才真正减少或近似历史 K/V。\n小结：系统管理的边界 # 系统管理能做的：显存碎片与预分配浪费 ✓、动态增长与 continuous batching ✓、KV-aware scheduling ✓、multi-turn/agent cache 生命周期 ✓、restoration/offloading/multi-tier storage ✓。\n系统管理做不到的：单个 token 的 KV 表示大小 ✗、KV Cache 随上下文线性增长 ✗、long output 中 decode cache 持续增长 ✗。\nSystem management improves utilization; compression reduces what must be stored or accessed. 这就自然引出下一部分——哪些 KV 值得保留。\n三、缓存压缩：哪些 KV 值得保留？ # 本章的发力点是 减少有效 $T$ 或 $\\text{bytes}$。系统管理回答「KV Cache 怎么放」，压缩回答「哪些 KV 保留、如何表示、何时压缩」。重点是 2025–2026 对「重要性」定义的修正，而不是旧的 SnapKV/PyramidKV。\n压缩方法的四条路线 # 路线 思路 代表方法 Token selection eviction / pruning SnapKV、DapQ、LaProx、KVP Representation compression quantization KIVI、KVQuant、TurboQuant Semantic compression chunk-level selection ChunkKV Decode-time compression long-output / reasoning-aware LongFlow、Moment-KV 旧基线：attention-based heuristic # 早期方法奠定基础，但大多依赖 attention pattern，且主要面向 long input：\n方法 核心原理 优势 局限 StreamingLLM sink tokens + recent window 稳定流式生成 语义选择弱 SnapKV observation window attention training-free，简单 input-side 与 decode-side 可能错位 PyramidKV layer-wise budget 考虑层间差异 仍依赖 attention pattern 关键铺垫：Attention score is useful, but not always sufficient. attention sink 说明「高 attention」不一定等于语义重要——开头 token 可能只是承担概率质量或位置稳定作用。这正是后面 DapQ 和 LaProx 对「重要性」定义做修正的出发点。\nposition-aware：DapQ（位置比内容更重要） # SnapKV 的问题：用 prompt 末尾的 observation window 近似未来 decode attention，但 input-side query 可能与 decode-time query 错位。\nDapQ 的关键是位置对齐：它认为构造 future query 时，位置比内容更重要。\n追加 pseudo tokens 赋予未来 decoding 位置 pseudo queries attend to prompt KV Top-K 选择 prompt KV DapQ training-free，比 observation window 更 decode-aligned；局限是主要压缩 prompt KV，pseudo position 是超参数。\noutput-aware：LaProx（attention score ≠ 真实输出贡献） # LaProx 指出 token 是否重要不只取决于 attention weight，还取决于 value 和 output projection。它把 eviction 重构为 output-aware 的矩阵乘法近似：\n$$\\text{Attention output} = A \\cdot V \\cdot W_O, \\qquad \\text{importance}_i \\propto \\|A_{:,i}\\| \\times \\|(V W_O)_i\\|$$即 token 对输出的真实贡献由 attention map $A$、value $V$、output projection $W_O$ 三者共同决定。这比纯 attention 启发式更接近真实输出贡献，代价是计算和实现复杂度更高，端到端收益依赖 kernel。\nlearning-based：KVP（从手工规则到学习式 policy） # 2026 年出现了 learned eviction：把 $(k_i, v_i, \\text{position}_i)$ 喂给一个轻量 policy / RL agent，输出 token ranking，保留 top-budget KV。KVP 学习的是 token 对未来 decoding 的 utility，不再依赖固定启发式，可适配不同 head/budget；代价是需要离线训练，泛化和部署成本更高。\n量化：从 scalar 到 vector quantization # 量化的背景是 KIVI / KVQuant 这类 activation-specific 低比特量化，处理 key/value 不对称、outlier、RoPE 效应。新进展 TurboQuant（QJL + PolarQuant）走向 online vector quantization：\n维度 scalar quant（KIVI / KVQuant） vector quant（TurboQuant） 压缩单位 单个标量 向量 关注点 低比特表示 near-optimal distortion、保持内积结构 token 可能丢 token 不丢 token 核心动机：attention score 取决于 $q^\\top k$，所以量化要尽量保持向量内积结构，而不是单纯追求每个元素的低比特。TurboQuant 很新，实际收益依赖 kernel 和 benchmark 范围。\nsemantic-aware：ChunkKV（以语义 chunk 为压缩单位） # token-level pruning 的反例：可能保住了关键词，却丢了条件/否定这种依赖语义连续性的内容。ChunkKV 以语义 chunk 而非孤立 token 为压缩单位：semantic chunking → chunk importance estimation → preserve coherent chunks，保留语义连续性和证据链。局限是 chunk 边界和粒度影响压缩率与质量。\nreasoning-aware：LongFlow（decode-time 压缩） # 要区分 long input 和 long output：前者是 prefill 后形成大 prompt KV，后者是 decode 中 CoT 持续增长。LongFlow（2026）面向 reasoning 的 long output，复用 attention 中间结果 $\\alpha_t$，用 $\\text{score}_i = \\|\\alpha_t^i v_i\\|$ 估计 token 重要性，把 attention + estimation + eviction 融合进一个 kernel。卖点是 zero-history + zero-cost estimation——复用已有 attention 中间量，几乎不额外算；局限是依赖 query 相似性，需要定制 kernel。\n动态压缩趋势与核心 trade-off # 动态压缩还有 RocketKV（coarse permanent eviction + dynamic sparse attention）、Moment-KV（momentum-based decode-time importance tracking）等方向。但压缩不是单纯追压缩率，共性 trade-off 有五条：\n共性 trade-off 含义 Memory saving ≠ latency improvement 省显存不等于端到端变快 Attention score ≠ 真实重要性 见 LaProx Long input ≠ Long output 是不同问题，见 LongFlow 高压缩率可能损害 retrieval/reasoning/code 任务敏感 真实收益依赖 kernel/scheduler/workload 离开系统谈压缩率无意义 压缩的灵魂是 memory-latency-quality-system 的平衡，而不是更高的压缩率。\n四、架构协同：从设计源头减少 KV Cache # 本章是全文重点。发力点是 $H_{kv}$ 与缓存表示，并进一步攻击有效 $T$——问题前移到模型设计阶段，而不是推理时再压缩。\n架构协同的路线全景 # 回到公式 $\\text{KV Cache Size} = 2 \\times L \\times T \\times H_{kv} \\times D_h \\times \\text{bytes}$，架构协同抓三个量：$H_{kv}$（减少 KV heads）、cached state dimension（latent 压维度）、有效 $T$（token-axis compressed KV）。\nMHA (H_kv=H_q) → MQA (H_kv=1) → GQA (1\u0026lt;H_kv\u0026lt;H_q) → MLA (latent state) → TransMLA (GQA→MLA 泛化) → Sparse/Compressed (Top-k / CSA / HCA) → MLRA / GDN (低秩可切分 / 固定状态) MHA → MQA → GQA：减少 KV heads # 最直接的降 KV Cache 方式是减少 KV head 数，这是「共享 K/V」的折中：\n架构 KV heads 特点 MHA $H_{kv} = H_q$ 表达能力强，cache 最大 MQA $H_{kv} = 1$ cache 最小，可能影响质量 GQA $1 \u003c H_{kv} \u003c H_q$ 质量与显存折中 MQA 保留多个 query heads，共享一组 K/V；GQA 在 MQA 的极端共享和 MHA 的完全独立之间折中。\nMLA：缓存 latent state，压维度 # MLA（DeepSeek-V2）把 KV Cache 优化前移到架构层：缓存的是低维 latent state，而不是完整 K/V。\nHidden state │ down projection ▼ 低维 latent state ← 缓存这个（而非完整 K/V） │ attention 时 up projection ▼ 还原出 K/V 用于 attention 从架构层减少 KV Cache，不需要 token eviction；但需要训练阶段支持，不是任意模型即插即用。\nTransMLA：让 MLA 成为更一般的注意力表达 # TransMLA（2025）建立 GQA 与 MLA-like latent attention 的结构联系，让已有 GQA/MHA 模型通过 transformation / reparameterization 向低 KV 架构迁移成为可能。代价是需要转换、训练或校准，不是无成本替换。\nDeepSeek V4 CSA/HCA：压序列长度（token-axis compressed KV）★ # 这是全文的技术核心。先讲清 MLA 与 V4 的边界：\nMLA：每个 token 的 cache entry 更窄（压 hidden dimension），但 cache length 仍随 $T$ 增长——一个历史 token 对应一个 cache entry。 DeepSeek V4 CSA/HCA：把历史 token group 压缩成更少的 compressed entries（压 sequence length），让远距离历史 cache entries 变少——直接攻击有效 $T$。 MLA（压维度） V4 CSA/HCA（压序列） 映射关系 1 个 token → 1 个 latent entry 1 个 token group → compressed entries cache length 仍随 $T$ 增长 远距离历史的 cache entries 更少 作用于公式 每个 entry 更窄（hidden dimension） 有效 $T$ 变小（sequence length） V4 的三条路径：\n路径 机制 用途 CSA（Compressed Sparse Attention） 温和块压缩 → compressed entries + indexer/top-k 选择 减少远距离历史的存储条目数和每步访问量 HCA（Heavily Compressed Attention） 更强的历史块压缩 → 在短 cache 上做 dense/global attention 用较低成本保留全局覆盖 SWA / local branch 保留最近窗口未压缩 KV 弥补压缩历史的 token-level 细节损失 报告中 1M context 场景下 V4-Pro 相对 V3.2 的单 token FLOPs 约 27%、KV cache 约 10% 是资料/公开报告口径，不是本文作者的复现实验。CSA/HCA 不是无损删除历史，而是用压缩历史换效率；实际 serving 收益依赖推理引擎如何分别管理 CSA/HCA/SWA 和 prefix cache。\nMLA 让每个 token 的 KV entry 更窄；CSA/HCA 让远距离历史的 cache entries 更少——两者作用于公式的不同维度。\n四条路线的边界：压维度 / 压序列 / 少访问 / 固定状态 # 这是必须讲清的核心区分（四条路线经常被混为一谈）：\n路线 改变什么 是否减少总存储 代表方法 MLA 压 hidden dimension（每个 entry 更窄） 是 DeepSeek-V2 CSA/HCA 压 sequence length（cache entries 更少） 是 DeepSeek V4 Top-k / page sparse 少访问（每步只读 $k$ 个，完整 KV 可能仍在） 不一定 Quest / MInference / LServe / NSA GDN 固定状态（不再为每个 token 存 KV） 是 Gated DeltaNet Top-k / page sparse attention 从 $O(T^2)$ → $O(Tk)$，$k$ 固定则线性。但要注意边界：完整 KV 仍在 GPU 时，sparse 主要降带宽和 kernel 计算，不一定降总显存；只有 full KV 移到 CPU/SSD、按需取 critical pages，才直接缓解 HBM 容量压力。\nGDN：从显式 KV Cache 到固定 recurrent state # softmax attention decode 时，cache $K_1, V_1 \\dots K_T, V_T$，每步 query $q_t$ 读全部历史 KV，memory $O(T)$、per-step read $O(T)$。\nGDN（Gated Delta Networks，ICLR 2025）走线性注意力路线，利用结合律把历史 key-value 外积累积为固定状态 $S_t$：\n$$S_t = \\text{gated\\_decay}(S_{t-1}) + \\text{delta\\_update}(k_t, v_t), \\qquad o_t = S_t^\\top q_t$$memory $\\approx O(d^2)$，不随 $T$ 线性增长。GDN 在简单累加上加入 gating（控制旧记忆衰减）和 delta rule（用在线更新修正 key→value 关联）；GDN-2 解耦 erase 和 write 的 channel-wise gate；FG²-GDN 把标量 $\\beta_t$ 改成 channel-wise，提高长上下文 associative recall。\nSoftmax attention decode Linear / GDN attention 历史 cache $K_1,V_1 \\dots K_T,V_T$ 维护固定状态 $S_t$ 每步 $q_t$ 读全部历史 KV $S_t = \\text{gated\\_decay}(S_{t-1}) + \\text{delta\\_update}(k_t,v_t)$，$o_t = S_t^\\top q_t$ 内存 $O(T)$，per-step read $O(T)$ $\\approx O(d^2)$，不随 $T$ 增长 MLRA：面向 tensor parallel 的低秩 latent attention # MLA 的瓶颈：tensor parallel 下 latent cache 不易切分。MLRA（2026）将 latent state 分成多个可分区的 low-rank branches，保留低 cache 的同时改善多 GPU 并行效率。\n这体现的是 architecture-system co-design：注意力架构不仅要省 cache，还要适合多 GPU serving。\n五、前沿趋势：任务、检索、验证、系统共同感知 # 前三部分是相对成熟的分类，这一部分是从单点方法走向 workload-aware 的新趋势。\nReasoning-aware：long output 成为新瓶颈 # 时代 workload Long input era 长 prompt → 短回答 Reasoning era 适中 prompt → 长 CoT / verification / reflection reasoning model 使 long output 成为新瓶颈：输入可能不长，但 CoT 很长，decoding 中 KV Cache 持续膨胀。代表工作如 LongFlow、Moment-KV、Hold Onto That Thought、DesireKV。核心是关注 decode-time KV 增长而非只压缩 prompt KV；难点是中间推理 token 是否可丢弃很难判断。\nRetrieval-aware：每步访问哪些 KV # KV 优化不一定只是压缩存储，也可以优化每步 decode 访问哪些 KV：current query → retrieve 相关 KV 子集 → attend to selected → 避免读 full cache。\nParisKV 的关键设计：drift-robust retrieval、GPU-native retrieval path、sink + recent + CPU backing store、4-bit reranking。KV 优化可以发生在访问侧，而不只是存储侧。\nVerifiable / Lossless：compressed draft + full verify # VeriCache（2026）走 speculative 风格：compressed KV cache 出 draft tokens，full KV cache（在 GPU 外）做 verify，accept或纠正。兼顾加速和输出一致性；代价是 verification 和 full cache 管理有额外开销。\nSystem-aware multi-tier serving # 未来 KV Cache 优化会把 compression、offloading、scheduler、kernel 和 workload 统一考虑。分层存储 GPU HBM ↔ CPU DRAM ↔ SSD/remote ↔ multi-GPU/distributed 成为标配，代表工作如 CacheFlow（3D-parallel restoration）、Kareto（tiered storage trade-off）、TokenCake/Continuum（agent cache lifecycle）、AsymCache。\n六、统一框架与未来工作 # 五层统一框架 # 把全文串起来，KV Cache 优化是一个 Workload-aware Policy（reasoning · agents · RAG · serving）统领下、三大支柱协同的问题：\nSystem Management — where to place \u0026amp; restore KV? — PagedAttention · CacheFlow · Kareto Cache Compression — which KV to keep / represent? — DapQ · LaProx · TurboQuant · LongFlow Architecture — how to reduce KV by design? — MLA · CSA/HCA · GDN · MLRA KV Cache 优化是算法、架构和系统共同作用的问题——没有任何单一方法是银弹。\n如何公平评测 KV Cache 方法 # 压缩率不是唯一指标。公平评测要看四个维度：\n维度 指标 Memory 峰值显存、cache size、可支持 batch/context Latency TTFT、TPOT、decode throughput、恢复开销 Quality LongBench、RULER、perplexity、retrieval/reasoning/code tasks System kernel 是否规则、scheduler 是否能利用、offloading 是否稳定 四个常见误区：\n误区 说明 Needle-in-a-Haystack 做得好 ≠ 长上下文强 单点检索不等于完整长上下文能力 高压缩率 ≠ 端到端加速 memory saving ≠ latency improvement attention score 高 ≠ 语义重要 见 LaProx 只看 prompt compression 忽略 reasoning long-output 的 decode-time growth 未来工作：从压缩率到 workload-aware co-design # Reasoning-aware cache policy Training-time cache compressibility Retrieval + compression hybrid KV memory Hardware/kernel-friendly sparse KV access Lossless or verifiable KV compression Realistic multi-turn/agent serving benchmarks Evaluation beyond Needle-in-a-Haystack（LongBench + RULER + serving metrics） The next stage is not just higher compression ratio, but memory-latency-quality-hardware-workload co-design.\n一句话总结 # KV Cache 优化不是单纯追求更高压缩率，而是在具体 workload 下共同平衡显存容量、显存带宽、延迟、输出质量、kernel 友好性和 serving 调度复杂度。\n参考 # 系统管理 # PagedAttention / vLLM — Kwon et al., SOSP 2023 Online Scheduling for LLM Inference with KV Cache Constraints — 2025 CacheFlow: 3D-Parallel KV Cache Restoration — 2026 Kareto: Multi-Objective Tiered Storage — 2026 TokenCake: KV-Cache-Centric Serving for Multi-Agent — 2025–2026 Continuum: Multi-Turn LLM Agent Scheduling with KV Cache TTL — 2026 缓存压缩 # StreamingLLM / Attention Sinks — 2023 SnapKV — 2024 PyramidKV — 2024 KIVI — 2024 KVQuant — NeurIPS 2024 DapQ — 2026 LaProx — 2026 KVP — 2026 TurboQuant — 2025/2026 ChunkKV — 2025 LongFlow — 2026 RocketKV — 2025 Moment-KV — 2026 架构协同 # One Write-Head / MQA — Shazeer, 2019 GQA — Ainslie et al., 2023 DeepSeek-V2 MLA — 2024 TransMLA — 2025 Native Sparse Attention / NSA — ACL 2025 Quest — ICML 2024 MInference — NeurIPS 2024 LServe — MLSys 2025 Gated Delta Networks — ICLR 2025 DeepSeek V4 CSA/HCA Compressed Attention — 2026（资料来源：DeepSeek V4 architecture report, 2026；CSA/HCA 解读参考 Sebastian Raschka (2026)、Together AI (2026)） MLRA — 2026 前沿趋势 # VeriCache — 2026 ParisKV — 2026 Hold Onto That Thought — 2026 DesireKV — 2026 评测与基础 # FlashAttention — Dao et al., NeurIPS 2022 LongBench RULER ","date":"2026年7月10日","externalUrl":null,"permalink":"/notes/llm-inference/kv-cache-optimization-survey/","section":"笔记","summary":"以 KV Cache Size = 2×L×T×H_kv×D_h×bytes 为锚点，系统梳理 2023–2026 年 KV Cache 优化全景：PagedAttention 与 KV-aware serving 的系统管理、DapQ/LaProx/TurboQuant/LongFlow 等压缩新方法、MHA→MLA→DeepSeek V4 CSA/HCA→GDN→MLRA 的架构演进，并讲清四条路线（压维度/压序列/少访问/固定状态）的边界。","title":"长上下文 LLM 推理中的 KV Cache 优化综述：系统管理、缓存压缩与架构协同","type":"notes"},{"content":" 原论文：Accelerating Large-Scale Reasoning Model Inference with Sparse Self-Speculative Decoding，Yilong Zhao 等，UC Berkeley / MIT / UW / NVIDIA 等，2025-12，arXiv:2512.01278。代码：github.com/sspec-project/SparseSpec。\n一句话摘要 # 针对推理语言模型（RLM）长输出推理的 memory-bound 瓶颈，用同一模型做 self-speculative decoding——verification 阶段的 full attention 顺便 dump 出 attention scores，Top-K 选出 critical tokens，作为接下来 k 步 draft 阶段的动态稀疏模式（PillarAttn），无需训练、无损、最高 2.13× 加速。\n问题动机：RLM 推理是 attention-bound # 推理语言模型（DeepSeek-R1、o1）动辄生成上万 token 的 CoT。自回归特性导致每生成一个 token 都要加载全部历史 KV-Cache，长输出把瓶颈从 compute-bound 推向 memory-bound。\nQwen3-8B / H100 / batch 128 / 输出 8192：平均每步加载 KV-Cache 耗 21 ms，占端到端 70%+ profiling（Fig.2）：compute 利用率 \u0026lt; 50%，memory bandwidth 打满；attention 占端到端 \u0026gt; 77% KV-Cache 总量随输出长度线性增长 关键洞察：MLP（GEMM）是 compute-bound、可被 batch 摊销权重加载；attention 是 memory-bound、各请求 KV 独立无法摊销。优化点在 attention 的 KV-Cache 访问。\n核心方法：稀疏自推测解码 # 为什么用 self-speculation + 稀疏注意力 # 传统推测解码需训练独立 draft 模型 → 数据工程复杂、对推理任务 OOD（EAGLE3 实测 acceptance \u0026lt; 2） self-speculation：同一模型当 draft + target，零训练 研究表明 KV-Cache 中 5% 的 token 就主导 attention 输出（Lin et al. 2025），近无损 把稀疏注意力当 draft model（只算 critical tokens），full attention 当 target 做 verification → 无损 PillarAttn（全文核心）— 复用 verify scores 的动态稀疏 attention # 两个设计要点：\n(a) 动态稀疏模式：上下文语义有空间局部性，以小步长（stride）周期性重新识别稀疏模式，stride 内固定 → 识别开销被摊销。\n(b) 零开销识别（overhead-free identification）——全文灵魂：\nstride 直接复用推测步数 $k$：每做 $k$ 步稀疏 draft → 做 1 步 full attention verification verification 本就要算全部 token 的 attention scores → 顺手 dump 出来，对 logits 和 log-sum-exp 缓存，rematerialize 出 scores GQA 下，scores 先在 $k$ 个 draft token、同组 query head 上取平均，再 Top-K 选 critical tokens 结果：识别 critical tokens 零额外计算/存储开销（对比 Quest 等需单独打分的方法） 第一性原理浓缩动机：「既然推测解码的 verify 阶段必然要做一次 full attention，那它算出来的 attention score 是不是正好可以白送给下一轮的稀疏 draft 用——这样动态稀疏的\u0026quot;识别开销\u0026quot;问题就不存在了？」\n速度理论模型（§3.2，核心公式） # 设 $M$=KV 总内存，$B$=batch，$k$=draft 长度，$\\alpha$=acceptance rate，$s$=稀疏比例。\nBaseline 单步：$T_{\\text{base}} = T_{\\text{GEMM}}(B) + T_{\\text{Attn}}(M)$\nSpec 每接受 token（一轮 k draft + 1 verify，产出 $k\\alpha+1$ token，含 bonus）：\n$$T_{\\text{spec}} = \\frac{k+1}{k\\alpha+1} T_{\\text{GEMM}}\\!\\left(\\tfrac{2k+1}{k+1}B\\right) + \\frac{1}{k\\alpha+1} T_{\\text{Attn}}\\!\\left(\\tfrac{ks+1}{k+1}M\\right)$$加速比 $\\eta = T_{\\text{base}} / T_{\\text{spec}}$。推导要点：\n项 系数 物理含义 GEMM $\\frac{k+1}{k\\alpha+1} \u003e 1$ spec 多做 draft GEMM，是代价（但 $B\u003c\\hat{B}$ 时近免费） Attention $\\frac{ks+1}{k\\alpha+1} \u003c 1$ draft 只读 $s$ 比例 KV，大幅省（典型省 80%） 极限分析：\nAttention 主导（长输出 RLM）：$\\eta \\to \\frac{k\\alpha+1}{ks+1}$，典型 $\\to 5\\times$ GEMM 主导（大 batch $B \\to \\hat{B}$）：$\\eta \\to \\frac{k\\alpha+1}{k+1} \u003c 1$，负优化 论文工作点（attention 占 77%）：理论 $\\approx 2.2\\times$，实测 2.13× ✅ 四大系统设计 # 挑战 设计 机制 负载波动 统一 batch scheduler 维护 k 个 bucket，贪心 bin-packing 让 draft/verify 混批，每步 GEMM 输入稳定在 $\\frac{2k+1}{k+1}B$ kernel 配置异构 fused sparse+full kernel persistent-kernel 风格，单 kernel 内 on-chip dispatch，比串行快 1.3× 显式同步 延迟验证 verify 请求 stall 一个 cycle，CPU 元数据清理与 GPU 计算重叠，省 20%+ 端到端 KV 利用不足 动态 KV-Cache manager 激进拉高并发 + chunk-wise 异步 offload 到 host（每步仅 18MB，带宽够用，cycle time +0.5%） 实验（§5） # Setup：Qwen3-1.7B/8B/14B，TP1/2/4，DGX-H100；AIME/OlympiadBench/LiveCodeBench，temp 0.6；$s=0.05, k=8$。\n关键结果：\n对比对象 加速倍数 vs vLLM（端到端） 最高 2.13×（Qwen3-1.7B/AIME） vs vLLM-NGram 最高 1.56× vs MagicDec 最高 1.36× vs TriForce 最高 1.76× vs EAGLE3（训练版） 持平或更高，且零训练 acceptance length：PillarAttn 6.16/8（α≈0.77），远超 NGram/EAGLE3（均 \u0026lt; 2） 消融：统一调度 / 动态 KV 管理 / 延迟验证 分别贡献 1.23× / 1.61× / 1.12× 局限与未来方向（§6） # 短上下文不适用：batch 打满 compute 后整体 compute-bound，方法无效 MoE 模型：本方法只动 attention 不动 FFN，可直接套用；MoE 每 expert 激活 token 少 → $\\hat{B}$ 上移 → 潜力更大 与 MTP/EAGLE3 组合成 hierarchical（类 TriForce）：MTP 当初级 draft → PillarAttn 次级 → full 终验 sparsity ratio / stride 固定为静态超参 → 自适应是空间 CPU offload 用 FIFO，与访问频率脱节（明确弱点） 改进方向（基于速度公式的理论评估） # 方向 杠杆 机理 上界 提高 α（Top-K 质量） 最大 直接放大 η α 0.75→0.95：η 1.53×→1.88× 自适应稀疏 s 中 进一步降 attention，但 α 会跟着降 需 trade-off Hierarchical (MTP+Pillar) 中 摊薄 verify 开销 等效 k↑ 访问频率感知 KV 管理 高（独立 paper 价值） 用 dump 的 score 做 LRU/LFU 替换，攻击 FIFO 弱点 跨节点池化场景 最值得深挖：输出长度重尾下的访问频率感知 KV 管理——PillarAttn 已免费产出\u0026quot;每个 page 访问热度\u0026quot;信号（dump 的 score），是天然零成本替换策略输入，目前被浪费。\n","date":"2026年6月20日","externalUrl":null,"permalink":"/notes/papers/sparsec-speculative-decoding/","section":"笔记","summary":"针对推理语言模型长输出的 memory-bound 瓶颈，用同一模型做 self-speculative decoding——verification 阶段顺手 dump 出 attention scores 做 Top-K，作为后续 draft 的动态稀疏模式，零训练、无损、最高 2.13× 加速。","title":"SparseSpec：加速推理模型的稀疏自推测解码","type":"notes"},{"content":"","date":"2026年6月20日","externalUrl":null,"permalink":"/tags/%E6%8E%A8%E6%B5%8B%E8%A7%A3%E7%A0%81/","section":"Tags","summary":"","title":"推测解码","type":"tags"},{"content":"","date":"2026年6月20日","externalUrl":null,"permalink":"/tags/%E6%8E%A8%E7%90%86%E5%8A%A0%E9%80%9F/","section":"Tags","summary":"","title":"推理加速","type":"tags"},{"content":"","date":"2026年6月20日","externalUrl":null,"permalink":"/tags/%E7%A8%80%E7%96%8F%E6%B3%A8%E6%84%8F%E5%8A%9B/","section":"Tags","summary":"","title":"稀疏注意力","type":"tags"},{"content":"","date":"2026年6月18日","externalUrl":null,"permalink":"/tags/cuda/","section":"Tags","summary":"","title":"CUDA","type":"tags"},{"content":"","date":"2026年6月18日","externalUrl":null,"permalink":"/tags/flashattention/","section":"Tags","summary":"","title":"FlashAttention","type":"tags"},{"content":"原论文：FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness，Tri Dao 等，Stanford / SUNY Buffalo，NeurIPS 2022，arXiv:2205.14135。\n一句话总结 # FlashAttention 通过 IO-aware 的 tiling 和重计算，在不改变 attention 数学定义的前提下，大幅减少 HBM 读写，实现 2-4× 加速和 5-20× 内存节省，是 attention 优化的基础构件。\n核心问题 # 标准 attention 的内存瓶颈：\nS = Q·K^T # 读写 N×N 矩阵到 HBM P = softmax(S) # 读写 N×N 矩阵到 HBM O = P·V # 读写 N×N 矩阵到 HBM 对于序列长度 N，需要 O(N²) 的 HBM 读写。GPU 计算速度 \u0026raquo; HBM 带宽，导致 memory-bound。\n核心方法 # IO-Aware Tiling # 将 Q, K, V 分块（tile），在 SRAM（高速缓存）中完成计算：\nfor tile_Q in Q: for tile_K, tile_V in K, V: # 在 SRAM 中计算局部 attention S_local = tile_Q · tile_K^T P_local = softmax(S_local) O_local += P_local · tile_V 关键优化 # 优化 说明 Tiling 分块加载到 SRAM，减少 HBM 访问 Online softmax 增量计算 softmax，避免存储完整 S Recomputation 反向传播时重计算 forward 中间值，不存储 内存复杂度 # 方法 内存 说明 标准 Attention O(N²) 存储 S, P FlashAttention O(N) 仅存储 O，中间值重计算 关键结果 # 指标 效果 加速比 2-4×（vs 标准 PyTorch） 内存节省 5-20× 精度 Exact：无近似，数学等价 序列长度 支持更长序列（内存不再是瓶颈） 后续版本 # 版本 改进 FlashAttention-2 更好的并行化，减少 non-matmul FLOPs FlashAttention-3 异步加载/计算，利用新硬件特性 FlashAttention-3 (Hopper) 针对 H100 的 Tensor Memory Accelerator 与稀疏注意力的关系 # FlashAttention 优化密集 attention 的 IO，不改变 O(N²) 计算复杂度。稀疏注意力减少计算量到 O(N) 或 O(N log N)。两者正交：\nFlashAttention + 稀疏模式：稀疏的 tile 计算 FlashAttention + 长序列：使 O(N²) 可接受的范围扩大 局限 # 计算复杂度未变：仍是 O(N²)，只是内存优化 序列长度上限：SRAM 容量限制 tile size 硬件依赖：针对特定 GPU 架构优化 实现要点 # Tile size 选择：平衡 SRAM 容量和并行度 Softmax 稳定性：online softmax 的数值稳定性 Kernel 融合：load/compute/store 流水线 个人理解 # FlashAttention 是 attention 优化的基础设施。它不改变 attention 的数学形式，只是更高效地实现。这使得它成为所有后续 attention 优化（包括稀疏注意力、量化 attention）的基础 kernel。理解 FlashAttention 是理解所有 attention 优化的前提。\n","date":"2026年6月18日","externalUrl":null,"permalink":"/notes/papers/flashattention/","section":"笔记","summary":"通过 IO-aware 的 tiling 和重计算，在不改变 attention 数学定义的前提下大幅减少 HBM 读写，实现 2-4× 加速与 5-20× 内存节省——attention 优化的基础构件。","title":"FlashAttention：IO 感知的快速精确注意力","type":"notes"},{"content":"","date":"2026年6月18日","externalUrl":null,"permalink":"/tags/gpu/","section":"Tags","summary":"","title":"GPU","type":"tags"},{"content":"","date":"2026年6月18日","externalUrl":null,"permalink":"/tags/%E6%B3%A8%E6%84%8F%E5%8A%9B%E4%BC%98%E5%8C%96/","section":"Tags","summary":"","title":"注意力优化","type":"tags"},{"content":"","date":"2026年6月15日","externalUrl":null,"permalink":"/tags/cs336/","section":"Tags","summary":"","title":"CS336","type":"tags"},{"content":"","date":"2026年6月15日","externalUrl":null,"permalink":"/tags/ddp/","section":"Tags","summary":"","title":"DDP","type":"tags"},{"content":"","date":"2026年6月15日","externalUrl":null,"permalink":"/tags/fsdp/","section":"Tags","summary":"","title":"FSDP","type":"tags"},{"content":"","date":"2026年6月15日","externalUrl":null,"permalink":"/tags/%E5%B9%B6%E8%A1%8C%E7%AD%96%E7%95%A5/","section":"Tags","summary":"","title":"并行策略","type":"tags"},{"content":"","date":"2026年6月15日","externalUrl":null,"permalink":"/tags/%E5%88%86%E5%B8%83%E5%BC%8F%E8%AE%AD%E7%BB%83/","section":"Tags","summary":"","title":"分布式训练","type":"tags"},{"content":" 基于 Stanford CS336 Lecture 7（Parallelism），Tatsunori Hashimoto \u0026amp; Percy Liang，Stanford Spring 2025。视频：YouTube，Slides：lecture_07.py。\n上周讲单 GPU 内部的并行（kernel fusion、tiling、shared memory）。本周扩展到多 GPU / 多机并行。核心主题：编排计算以避免数据传输瓶颈。无论单 GPU 内部还是多 GPU 之间，根本问题都是——算力离数据太远。\n广义的存储层级：单 GPU 内 L1/shared memory（最快）→ 单 GPU 内 HBM → 单节点多 GPU NVLink/NVSwitch → 多节点 InfiniBand/Ethernet（最慢）。\n分布式通信 / 计算的基础构件 # 为什么需要多 GPU？ # 内存不够：参数量 + 优化器状态 + 梯度 + 激活值，单卡放不下 算力不够：想要更多 FLOPs 缩短训练时间 集合通信操作（Collective Operations） # 基本概念：Rank = 一个设备/GPU（如 0,1,2,3）；World Size = 总设备数（如 4）。\n类别 操作 说明 基础 Broadcast rank 0 复制到所有 rank Scatter rank 0 切分后分发，每人拿 1/world_size Gather 所有 rank 汇集到 rank 0（scatter 的逆） Reduce 所有 rank 聚合（SUM/MAX/MIN）到 rank 0 主力 All-gather 每人把自己的 shard 广播给所有人，每人最终持有完整拼接 Reduce-scatter 先 reduce 再 scatter，每人只拿到结果的 1/world_size All-reduce 先 reduce 再 broadcast 给所有人，每人拿完整聚合结果 特殊 All-to-all 每人对每人发一份数据（MoE 路由的关键操作），等价于转置 关键恒等式：\n$$\\text{all-reduce} = \\text{reduce-scatter} + \\text{all-gather}$$All-reduce 一步到位： rank0: [0,1,2,3] → [6, 10, 14, 18] rank1: [1,2,3,4] → [6, 10, 14, 18] rank2: [2,3,4,5] → [6, 10, 14, 18] rank3: [3,4,5,6] → [6, 10, 14, 18] 等价于 reduce-scatter + all-gather： Reduce-scatter: rank0=[6], rank1=[10], rank2=[14], rank3=[18] All-gather: 所有人=[6,10,14,18] 带宽效率上两者等价，但拆成两步给了 FSDP/ZeRO 灵活性的空间。\n命名记忆法：Reduce = 聚合（求和/取最大/取最小）；Scatter 是 Gather 的逆；All = 结果发给所有设备。\n硬件互联 # 层级 介质 带宽（典型值） 物理拓扑 单节点多 GPU NVLink 5.0 + NVSwitch ~1.8 TB/s（B200） 全互联（all-to-all） 多节点（同 pod） InfiniBand ~0.05 TB/s 全互联 ≤256 GPU 跨 pod / 跨数据中心 Ethernet ~200 MB/s ~ 0.05 TB/s 叶脊交换机 绕过 CPU 的技术：\nRDMA（Remote Direct Memory Access）：GPU 直接读写远端 GPU 内存，不经 CPU RoCE（RDMA over Converged Ethernet）：让 Ethernet 也支持 RDMA（Meta 在用） InfiniBand 原生支持 RDMA，标准 Ethernet 不支持 GPU vs TPU 网络拓扑差异：GPU 节点内 8 GPU 全互联（NVSwitch），≤256 GPU 内任意通信都快；TPU 是 3D 环面网格（Toroidal Mesh），芯片只和邻居连接，可轻松扩展但只能邻居通信。对于集合通信，两者在理论上效率相同。\nNCCL 与 PyTorch Distributed # NCCL（NVIDIA Collective Communication Library）把集合通信操作翻译成 GPU 间传输的低层数据包：检测硬件拓扑 → 优化通信路径 → 启动 GPU kernel 做收发。\nimport torch.distributed as dist # 初始化 os.environ[\u0026#34;MASTER_ADDR\u0026#34;] = \u0026#34;localhost\u0026#34; os.environ[\u0026#34;MASTER_PORT\u0026#34;] = \u0026#34;15623\u0026#34; dist.init_process_group(\u0026#34;nccl\u0026#34;, rank=rank, world_size=world_size) # 核心操作 dist.all_reduce(tensor, op=dist.ReduceOp.SUM) # 修改 in-place dist.reduce_scatter_tensor(output, input, op=dist.ReduceOp.SUM) dist.all_gather_into_tensor(output_tensor, input_tensor) dist.broadcast(tensor, src=0) # 清理 dist.destroy_process_group() 后端：gloo（CPU）、nccl（GPU）。\nAll-reduce 的有效带宽公式：$\\text{bandwidth} \\approx \\frac{2 \\times \\text{size\\_bytes}}{\\text{duration}}$，与 world_size 无关，与拓扑无关。\n分布式训练算法 # Data Parallelism（DDP） # 每张卡持有完整模型副本，把 batch 切分到各卡。\n# 每个 rank 拿自己的数据切片 data = data[rank * local_bs : (rank+1) * local_bs] for step in range(num_steps): loss = forward(data, params) loss.backward() # 唯一区别于单卡训练的地方 for param in params: dist.all_reduce(param.grad, op=dist.ReduceOp.AVG) # 注意是 AVG optimizer.step() 维度 评分 说明 计算扩展 ✅ 好 每卡拿到 B/world_size 样本，batch 够大就能打满算力 通信开销 ⚠️ 中等 每步 all-reduce 2×参数量，batch 大可以隐藏 内存扩展 ❌ 差 每张卡都要存完整模型 + 优化器状态 使用 ReduceOp.AVG 而非 SUM：AVG 自动除以 world_size。\n为什么 DDP 的内存问题这么严重？ # 以 AdamW 优化器为例，一个参数需要存储：\n组件 精度 字节数 参数（weights） BF16 2 梯度（gradients） BF16 2 主权重（master weights） FP32 4 Adam m（一阶矩） FP32 4 Adam v（二阶矩） FP32 4 总计 16 bytes/param 内存大头是优化器状态（m + v + master weights = 12 bytes），占 75%。这就是为什么需要 ZeRO / FSDP。\nFSDP / ZeRO（Fully Sharded Data Parallel） # 渐进式内存节省：\n阶段 切分内容 内存节省（相对 DDP） ZeRO-1 优化器状态 ~4× ZeRO-2 优化器状态 + 梯度 ~8× ZeRO-3 优化器状态 + 梯度 + 参数 线性（world_size 倍） 数据流（ZeRO-3 / FSDP）：\nForward: all-gather params → 拼出完整参数（临时）→ 前向计算 → 释放完整参数 Backward: 反向计算 → 每卡算出完整梯度 reduce-scatter grads → 每卡只保留自己的梯度 shard Optimizer: 每卡只对自己的参数 shard 做 optimizer.step() 每卡只存自己的优化器状态 核心操作对：all-gather（forward 拼参数）+ reduce-scatter（backward 收梯度）。代价：每个 layer 都要做 all-gather 和 reduce-scatter，有同步 barrier，比 DDP 更多通信量和同步点。这就是开头强调 all-reduce = reduce-scatter + all-gather 的原因——DDP 用 all-reduce 一步完成，FSDP 拆成两步，换来了内存线性扩展。\nTensor Parallelism # 横切——把每层的权重矩阵沿列/行方向切开，每张卡只保留矩阵的一「列」。\n# 每张卡用自己的参数分片做局部计算 x = x @ params[layer] # (B, local_dim) # All-gather 把各卡的激活值拼回完整维度 activations = [torch.empty(B, local_dim) for _ in range(world_size)] dist.all_gather(tensor_list=activations, tensor=x) x = torch.cat(activations, dim=1) # (B, num_dim) 维度 评分 说明 计算扩展 ✅ 好 算力随 GPU 数线性增长 通信开销 ❌ 高 每层都要 all-gather，需要极快互联 内存扩展 ✅ 线性 参数/激活都能切分 对 batch size 影响 ✅ 无 唯一不消耗 batch size 的并行 实践规则：TP 只在单节点内做（NVLink 带宽够），通常 TP=8。\nPipeline Parallelism # 纵切——把模型的不同层放到不同 GPU 上，数据像流水线一样穿过各 GPU。\nPipeline Bubble 问题：GPU 间有空闲时间。解决方案：Micro-batches——把一个大 batch 切成多个 micro-batch 交错执行减少空闲。micro-batch 越多 bubble 越小，但需要更多内存存中间激活值。核心是用 P2P 通信（dist.send / dist.recv），不需要集合通信。\n维度 评分 说明 计算扩展 ⚠️ 中等 有 pipeline bubble 通信开销 ✅ 低 点对点传输，带宽要求低 内存扩展 ✅ 线性 每卡只存一个 stage 的参数 对 batch size 影响 ❌ 消耗 micro-batch 数消耗有效 batch size 工程复杂度 ❌ 高 需要精细调度 micro-batch Sequence Parallelism \u0026amp; Activation Memory # 即使用了 TP，激活内存中仍有一些项没法被 TP 切分——LayerNorm、Dropout 等 point-wise 操作。\n激活内存公式（Transformers 单层）：\n$$\\text{activation memory} = SBH \\times 34 + \\frac{5AS^2B}{H}$$ 左边（34 SBH）：MLP 和点操作的激活（取决于 H） 右边（5AS²B / H）：Attention softmax 的中间结果（和 S² 成正比） 用了 TP（切分到 T 个设备）后：\n$$\\text{memory after TP} = \\frac{SBH \\times 34}{T} + \\frac{5AS^2B}{H}_{\\text{被 FlashAttention 消掉}} + SBH \\times 10$$ 第二项被 FlashAttention 消掉（recomputation） 第三项 SBH × 10 是 LayerNorm、Dropout 等不受 TP 影响的残留 Sequence Parallelism 沿序列维度切分这些 point-wise 操作，最终收敛到 $\\text{minimal activation memory} \\approx \\frac{SBH \\times 34}{T}$，这是 TP + FlashAttention + Sequence Parallelism 的极限。\n其他并行策略 # 策略 思路 适用场景 Ring Attention / Context Parallel 切分长序列的 attention，KV 在设备间循环传递 超长上下文训练 Expert Parallelism MoE 中不同 expert 放不同设备，用 all-to-all 路由 MoE 模型 组合使用 # 并行策略对比总表 # 策略 切分维度 通信开销 内存扩展 消耗 batch size 适用网络 DDP Batch 每步 1 次 all-reduce ❌ 无 否 任意 FSDP (ZeRO-3) Batch+参数 每层 all-gather + reduce-scatter ✅ 线性 否 IB/Ethernet Tensor Parallel Width 每层 all-gather ✅ 线性 否（唯一！） NVLink Pipeline Parallel Depth 低（P2P） ✅ 线性 是 IB/Ethernet Sequence Parallel Sequence all-gather + reduce-scatter 辅助 否 NVLink 三种有限资源 # Memory：决定模型能不能跑 Bandwidth + Compute：决定跑得快不快 Batch Size：影响通信隐藏效率——batch 太小则通信 overhead 占比高 实践法则（Rule of Thumb） # 第一步：让模型放进内存（硬约束） → Tensor Parallelism：先铺满单节点内 GPU（通常 TP=8），利用 NVLink 的高带宽 → FSDP (ZeRO-3) 或 Pipeline Parallelism：跨节点扩展直到模型能装下 第二步：用剩余的 GPU 做 Data Parallelism 放大吞吐 → DP 带宽要求低，最灵活 第三步（可选）：如果 batch size 太小 → Gradient Accumulation：多步再同步一次梯度，等价于增大有效 batch size 带宽从高到低的并行策略层：\nTP (NVLink) → CP (NVLink) → PP (IB) → DP (IB/Ethernet) 高带宽需求 ←─────────────────────────────→ 低带宽容忍 真实案例 # 模型 并行策略 细节 Megatron-LM (530B) TP=8 + PP 递增 + DP 调整 1.7B→1T 参数，TP 在 8 封顶，大模型加 PP DeepSeek-V3 16-way PP + 64-way Expert Parallel + ZeRO-1 DP 用 Expert Parallel 替代 TP Llama 3 (405B) TP=8 + CP（长上下文）+ PP + DP 严格按 TP→CP→PP→DP 带宽递减顺序 Gemma 2 (TPU) ZeRO-3 + Model Parallelism TPU 3D Torus 允许更大的 Model Parallel 范围 Llama 3 的血泪教训：训练中断共 466 次，其中 148 次 来自 GPU 硬件故障（占 30%）；计划外维护导致 32 次 中断；静默数据损坏（silent data corruption）比显式故障更可怕——GPU 可能输出错误的数值而不报任何错误。\n关键恒等式速查 # 恒等式 含义 all-reduce = reduce-scatter + all-gather DDP 一步 = FSDP 两步，带宽等价 有效 BW ≈ 2×size/duration all-reduce 带宽公式（与 world_size 无关） memory/param = 16 bytes (AdamW) 参数 2 + 梯度 2 + master 4 + m4 + v4 activation ≈ SBH×34/T TP + FlashAttn + SeqParallel 后的极限 参考 # 课程网站：cs336.stanford.edu Slides 源码：github.com/stanford-cs336/lectures 视频：YouTube ","date":"2026年6月15日","externalUrl":null,"permalink":"/notes/systems/cs336-distributed-parallelism/","section":"笔记","summary":"从单 GPU 扩展到多 GPU/多机并行：集合通信原语（all-reduce = reduce-scatter + all-gather）、NVLink/InfiniBand 互联，以及 DDP、FSDP/ZeRO、Tensor/Pipeline/Sequence Parallelism 的取舍与实践法则。","title":"分布式训练并行策略：CS336 Lecture 7 笔记","type":"notes"},{"content":" 基于 Stanford CS336 Spring 2026 Lecture 6（Kernels, Triton, XLA）整理。官方材料：lecture_06.py。\nLecture 5 从硬件角度建立 GPU 直觉，Lecture 6 从实践角度进入 kernel 编程和 profiling。三条线：\nGPU 硬件回顾 + 编程模型与硬件的交互（warp/occupancy/bank conflict/coalescing） 性能分析：benchmarking + profiling，理解瓶颈在哪 手写 kernel：用 Triton 实现 GeLU、softmax、row sum、matmul+ReLU 硬件参数对比 # Accelerator A100 H100 B200 # SMs 108 132 148 Register (per SM) 256 KB 256 KB 256 KB L1 + shared (per SM) 192 KB 256 KB 256 KB L2 cache 40 MB 50 MB 96-126 MB HBM 80 GB 80 GB 192 GB Register bandwidth ~116 TB/s ~401 TB/s ~447 TB/s L1 + shared bandwidth ~19 TB/s ~33 TB/s ~19 TB/s L2 bandwidth ~5-8 TB/s ~12 TB/s ~9 TB/s HBM bandwidth 2 TB/s 3.35 TB/s 8 TB/s B200 还有一个对 programmer 不可见的 tensor memory（TMEM），位于 register 和 shared memory 之间，专门用于 Tensor Core。\nProgramming model：\nGPU kernel └── Grid (所有 block) └── Thread block / CTA (一组 threads，共享 shared memory) └── Thread (最小执行单元，私有 register) Thread：执行一小部分数据的代码 Thread block / CTA (Concurrent Thread Array)：一组共享 shared memory 的 threads Grid：所有 thread blocks 的集合 H100/B200 还有 thread block clusters，支持 distributed shared memory。\n编程模型与硬件的交互 # 核心思想：程序模型提供了正确的抽象，但性能极度依赖对硬件的理解。\nWarps # 一个 thread block 内的 threads 被分为 warps，每个 warp 32 个 threads。 所有 warp 内 threads 在 SM 上以 lockstep 执行同一条指令。 Control divergence：如果同一个 warp 内不同线程需要走不同分支，GPU 会串行执行两个路径（先 mask 掉一部分线程执行 path A，再 mask 掉另一部分执行 path B），总时间 ≈ 两个分支时间之和。 Zero-cost warp switching：SM 可以在多个 warp 之间切换（例如当某个 warp 等待 HBM 读写时），切换开销几乎为零。 Warp Occupancy # 每个 thread 可使用 0-255 个 registers。 threads 用越多 registers，SM 上能同时调度的 threads 就越少（低 occupancy）。 Low occupancy 不一定是坏事：如果每个 thread 做更多工作（thread coarsening），总吞吐可能更高。 示例计算：\nthread block = 128 threads 每个 thread 用 160 registers SM 最多 65536 registers → 每 SM 最多 65536 / (128 × 160) = 3 个 blocks warps = 3 × 128 / 32 = 12 occupancy = 12 / 64 = 18.75% 较高的 register 使用降低 occupancy，但可能意味着每个 thread 做了更多有用的工作（例如处理更多元素），最终可能更快。\nBank Conflicts (Shared Memory) # Shared memory 分为 32 个 banks，每个 4 字节宽。每个周期，每个 bank 只能被一个 thread 访问（除非访问完全相同的位置）。\n如果多个 threads 访问同一个 bank（不同位置），访问被串行化 → bank conflict。 最坏情况：矩阵每行跨所有 banks，32 个 threads 访问第一列 → 32-way bank conflict。 解决方式：swizzling，通过重新排列 shared memory 布局（如 row xor col）来避免冲突。 Memory Coalescing (HBM) # 当一个 warp 的 32 个 threads 访问 HBM 时，内存访问会合并为 128 字节的缓存行（cache line）事务。\n最佳情况：全部 coalesced，所有 threads 访问同一 cache line 内的连续地址（32 threads × 4 bytes = 128 bytes）。\nBlock Occupancy 与 Wave Quantization # Thread blocks 按 wave 调度到 SM 上。B200 有 148 个 SMs，如果 launch 160 个 thread blocks，第一 wave 148 个，第二 wave 12 个。Wave quantization 问题：最后一 wave 的 block 太少，部分 SM 空闲。解决：让 thread blocks 数量尽量被 SM 数量整除。\nBenchmarking 与 Profiling # 方法论三步循环：\n1. Benchmark 和 profile 你的代码 2. 做优化 3. 再次 benchmark 和 profile Benchmarking # 测量端到端的 wall-clock time。为什么要用 CUDA events 而不是 Python 的 time.time()？\nstart_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) start_event.record() run() end_event.record() torch.cuda.synchronize() time_ms = start_event.elapsed_time(end_event) CUDA events 直接在 GPU 侧记录时间，避免了 CPU 调度开销和 Python 解释器干扰。三个关键步骤：Warmup（不计入编译/JIT 开销）、同步（torch.cuda.synchronize()）、多次运行取平均。\nProfiling # Benchmarking 只告诉你\u0026quot;慢不慢\u0026quot;，profiling 告诉你\u0026quot;慢在哪\u0026quot;。用 torch.profiler：\nwith torch.profiler.profile(activities=[ProfilerActivity.CUDA], experimental_config=torch._C._profiler._ExperimentalConfig(verbose=True)) as prof: run() torch.cuda.synchronize() table = prof.key_averages().table(sort_by=\u0026#34;cuda_time_total\u0026#34;, row_limit=10) 关键发现：\nadd(dim=2048)：少量简单 kernel，memory-bound。 matmul(dim=2048)：ampere_sgemm* 系列 GEMM kernel，compute-bound。 matmul(dim=128)：小矩阵时 kernel 名字不同，可能不是 GEMM 而是更小的 kernel template。 Naive vs Builtin vs Compiled GeLU # 用 GeLU 演示 kernel fusion 的重要性。\nNaive GeLU：从公式拼出，每一步产生独立 kernel 0.5 * x * (1 + torch.tanh(0.79788456 * (x + 0.044715 * x * x * x))) 多个元素操作 → 多个 kernel launches → 多次 HBM 读写。 Builtin GeLU：torch.nn.functional.gelu(x, approximate=\u0026quot;tanh\u0026quot;)，PyTorch 内置的 fused 实现。 Compiled GeLU：torch.compile(naive_gelu)，PyTorch JIT 自动分析和融合，生成单个 Triton kernel。 实现 kernel 数量 HBM 读写 速度 naive 多个 多次 read/write 最慢 builtin 单个 fused 一次 read + 一次 write 快速 compiled 单个 Triton kernel 一次 read + 一次 write 接近 builtin 核心 insight：未融合的 elementwise chain 中，每个中间结果都要写入 HBM 再读回。Fusion 后中间结果留在 register 中，大幅减少 HBM traffic。\nTriton Kernel 编程 # Triton（OpenAI 开发）的编程模型与 CUDA 不同：\nCUDA: 程序员指定每个 thread 做什么 Triton: 程序员指定每个 thread block 做什么 思维框架：load 数据到 shared memory → 在片上操作 → 写回 global memory。\nTriton GeLU（Elementwise） # @triton.jit def triton_gelu_kernel(x_ptr, y_ptr, num_elements, BLOCK_SIZE: tl.constexpr): pid = tl.program_id(axis=0) # 当前 block 的 ID start = pid * BLOCK_SIZE # 该 block 的起始偏移 offsets = start + tl.arange(0, BLOCK_SIZE) # 该 block 的索引范围 mask = offsets \u0026lt; num_elements # 边界检查 x = tl.load(x_ptr + offsets, mask=mask) # 计算 GeLU a = 0.79788456 * (x + 0.044715 * x * x * x) exp = tl.exp(2 * a) tanh = (exp - 1) / (exp + 1) y = 0.5 * x * (1 + tanh) tl.store(y_ptr + offsets, y, mask=mask) 关键点：tl.program_id(axis=0) 获取 block 索引；tl.arange(0, BLOCK_SIZE) 生成线程索引范围；mask 用于边界检查；声明 BLOCK_SIZE: tl.constexpr 使编译器编译时确定 block size。\nTriton Softmax（单行 reduction） # @triton.jit def triton_softmax_kernel(x_ptr, y_ptr, x_row_stride, y_row_stride, num_cols, BLOCK_SIZE: tl.constexpr): assert num_cols \u0026lt;= BLOCK_SIZE # 一行必须能放入一个 block row_idx = tl.program_id(0) # 每行由一个 block 处理 col_offsets = tl.arange(0, BLOCK_SIZE) x_start_ptr = x_ptr + row_idx * x_row_stride x_row = tl.load(x_start_ptr + col_offsets, mask=col_offsets \u0026lt; num_cols, other=float(\u0026#34;-inf\u0026#34;)) # softmax: x - max → exp → sum → normalize x_row = x_row - tl.max(x_row, axis=0) numerator = tl.exp(x_row) denominator = tl.sum(numerator, axis=0) y_row = numerator / denominator y_start_ptr = y_ptr + row_idx * y_row_stride tl.store(y_start_ptr + col_offsets, y_row, mask=col_offsets \u0026lt; num_cols) 每行一个 block，tl.max、tl.sum 是 Triton 的 reduction 操作符，内部处理跨线程通信。相比 naive 实现，Triton 版本在 shared memory 中完成所有操作，只有一次 read 和一次 write。\nTriton Row Sum（大行需要循环的 reduction） # 如果一行比 BLOCK_SIZE 大，需要 tiling：\n@triton.jit def row_sum_kernel(x_ptr, out_ptr, N, BLOCK_SIZE: tl.constexpr): row = tl.program_id(0) acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32) # 每个线程的累加器 for start in range(0, N, BLOCK_SIZE): # 循环处理 tiles cols = start + tl.arange(0, BLOCK_SIZE) mask = cols \u0026lt; N x = tl.load(x_ptr + row * N + cols, mask=mask, other=0.0) acc += x result = tl.sum(acc, axis=0) # 最终 reduction tl.store(out_ptr + row, result) 这是 Triton 的 \u0026ldquo;baby tiling\u0026rdquo; 模式：当行太大无法放入单个 block 时，每个 block 依次处理多个 tiles，在 acc 中累加，最后用 tl.sum 做跨线程 reduction。\nTriton Matmul + ReLU（Tiling + Kernel Fusion） # Naive matmul 的问题：每次读 $A[m,k]$ 和 $B[k,n]$ 从 HBM，计算后写 $C$ 到 HBM → $M \\times K \\times N$ 次 HBM 读，总共只做 $O(1)$ arithmetic intensity。\nTiling 方案：\n1. 把 C 分成输出 tile（如 64×64），每个 block 负责一个 tile 2. 每次从 HBM 加载一对 A tile（如 64×32）和 B tile（如 32×64）到 shared memory 3. 在片上做这小块 matmul，累加到部分和 4. 重复直到覆盖整个 K 维度 5. 写回 HBM Arithmetic intensity ≈ O(tile_size)，远高于 naive 的 O(1) @triton.jit def matmul_relu_kernel( a_ptr, b_ptr, c_ptr, M, N, K, strides..., BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, ): pid_m = tl.program_id(0) pid_n = tl.program_id(1) indices_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) indices_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) indices_k = tl.arange(0, BLOCK_K) a_ptrs = a_ptr + indices_m[:, None] * stride_am + indices_k[None, :] * stride_ak b_ptrs = b_ptr + indices_k[:, None] * stride_bk + indices_n[None, :] * stride_bn acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) for k in range(0, K, BLOCK_K): a = tl.load(a_ptrs, mask=..., other=0.0) b = tl.load(b_ptrs, mask=..., other=0.0) acc += tl.dot(a, b) # GPU Tensor Core 加速 a_ptrs += BLOCK_K * stride_ak # 推进到下一块 b_ptrs += BLOCK_K * stride_bk acc = tl.maximum(acc, 0.0) # fused ReLU c_ptrs = ... tl.store(c_ptrs, acc, mask=...) Kernel fusion 的优势：在写回 HBM 之前，ReLU（elementwise）在 register 中完成，不需要额外的 HBM 往返。\n理解 CUDA Kernel 名字 # profiling 输出的 kernel 名字里有大量信息：\ncutlass3x_sm100_simt_sgemm_f32_f32_f32_f32_f32_64x64x16_1x1x1_3_nnn_align1_... 片段 含义 sm100 Blackwell B200 架构 simt SIMT 路径（非 Tensor Core 路径） sgemm Single precision GEMM f32 float32 精度 64x64x16 tile shape: 输出 C tile 64×64, K tile 16 nnn 矩阵转置标记（n = not transposed） ampere_sgemm_128x128_* 则是用于 Ampere A100 的 GEMM kernel，128×128 是 tile size。\n从 Triton 到 PTX # Triton 编译生成的 PTX（Parallel Thread Execution）是 GPU 的中间表示/汇编层。Triton GeLU 生成的 PTX 包含：ld.global.*（从 global memory 读取）、st.global.*（写入 global memory）、%ctaid.x（block index）、%tid.x（thread index）。Triton 自动做了 thread coarsening——一个 thread 处理 8 个元素。\n术语表 # 术语 含义 CTA Concurrent Thread Array，与 thread block 同义 Warp GPU 调度单位，32 个线程 SIMT Single Instruction Multiple Threads Occupancy SM 上活跃 warp 占理论上限的比例 Bank conflict 多个线程访问 shared memory 同一 bank 导致串行化 Memory coalescing warp 内线程访问连续 HBM 地址，合并为单个事务 Wave quantization 最后一批 thread block 数量少于 SM 数，造成部分 SM 空闲 Swizzling 重新排列 shared memory 布局以规避 bank conflict Thread coarsening 每个线程处理多个元素以提高利用率 Kernel fusion 合并多个 kernel 以减少 HBM 读写 PTX Parallel Thread Execution，GPU 汇编级 IR Triton OpenAI 的 GPU 编程语言，以 thread block 为单位编程 CUDA events GPU 侧计时器，用于精确的 GPU 时间测量 cutlass NVIDIA 的 CUDA 线性代数模板库 总结 # - 了解编程模型（PyTorch, Triton, PTX）以保证正确性 - 理解硬件（SMs, warps, occupancy, bank conflicts 等）以优化性能 - Benchmark 以了解 scaling 行为 - Profile 以查看在运行什么、跑多久 - Triton 以 thread block 为单位思考（load到shared memory → 操作/fusion → 写回HBM） - 示例：GeLU (elementwise), softmax (row-wise), row sum (tiling), matmul (tiling + fusion) ","date":"2026年6月14日","externalUrl":null,"permalink":"/notes/triton/cs336-kernels-triton/","section":"笔记","summary":"从硬件抽象进入 kernel 编程实践：warp/occupancy/bank conflict/coalescing 如何映射到性能，benchmarking 与 profiling 方法论，以及用 Triton 实现 GeLU、softmax、row sum、matmul+ReLU 四个 kernel。","title":"Kernels、Triton 与 Profiling：CS336 Lecture 6 笔记","type":"notes"},{"content":"","date":"2026年6月14日","externalUrl":null,"permalink":"/tags/profiling/","section":"Tags","summary":"","title":"Profiling","type":"tags"},{"content":"","date":"2026年6月14日","externalUrl":null,"permalink":"/tags/triton/","section":"Tags","summary":"","title":"Triton","type":"tags"},{"content":"","date":"2026年6月14日","externalUrl":null,"permalink":"/tags/%E6%80%A7%E8%83%BD%E4%BC%98%E5%8C%96/","section":"Tags","summary":"","title":"性能优化","type":"tags"},{"content":" 基于 Stanford CS336 Spring 2026 Lecture 5（GPUs, TPUs, and Efficient Attention）整理。官方材料：lecture_05.pdf。\n本讲从模型结构转入系统视角，目标是理解：为什么现代 LLM 训练和推理必须依赖 GPU/TPU，以及为什么同样的数学计算在不同 kernel、不同内存访问模式下性能可能相差巨大。\n主线 # 本讲可以概括为三条线：\nGPU/TPU 为什么适合 LLM GPU 的执行模型和内存层次 如何用性能模型解释 FlashAttention 这类 IO-aware 算法 核心不是学习 CUDA 语法，而是建立系统直觉：\n快的深度学习程序 = 足够多的并行性 + 高效的数据复用 + 避免昂贵的显存访问 对于 LLM，大量计算集中在矩阵乘法、attention、MLP、normalization 和 elementwise 操作。矩阵乘法能很好利用 Tensor Cores，但 attention 和许多小算子容易受内存带宽、kernel launch、数据搬运影响。\n为什么 LLM scaling 依赖 GPU scaling # LLM 的能力提升通常来自：更多参数、更多训练 token、更多训练 FLOPs、更高硬件利用率、更强并行化能力。\n过去 CPU 单核性能增长依赖频率提升、工艺缩放和 Dennard scaling。Dennard scaling 放缓后，单核性能无法继续高速增长。现代深度学习 scaling 更依赖 parallel scaling——通过大量并行计算单元提升吞吐。\nGPU 的优势在于：\nCPU: 优化少量复杂线程的低延迟执行 GPU: 优化海量简单线程的高吞吐执行 这和 LLM 的计算模式高度匹配，因为 Transformer 中大部分算子都可以表达为大批量矩阵或张量操作。\nCPU vs GPU：latency-oriented 与 throughput-oriented # CPU 更关注：少量线程、复杂控制流、强分支预测、大缓存、低延迟响应。\nGPU 更关注：大量线程、简单控制逻辑、高并行吞吐、更高算术单元密度、用线程切换隐藏内存延迟。\nGPU 并不是在每个线程上都比 CPU 快，而是能同时运行大量线程，让整体吞吐非常高。\nCPU 像少数很强的工人，每个人能处理复杂任务。 GPU 像大量专门工人，每个人做简单操作，但总吞吐极高。 在深度学习中，我们通常不关心单个 token 或单个元素的延迟，而关心整个 batch、整个矩阵乘法、整个训练 step 的吞吐。\nGPU 的基本硬件结构 # GPU 由多个 Streaming Multiprocessors（SM）组成：\nGPU └── many SMs ├── CUDA cores / SPs ├── Tensor Cores ├── registers ├── shared memory / L1 cache └── warp schedulers 术语 含义 SM Streaming Multiprocessor，执行 thread block 的主要计算单元 SP / CUDA core 执行普通标量或向量浮点运算的计算单元 Tensor Core 专门加速矩阵乘法的硬件单元 register 每个线程私有的最快存储 shared memory 一个 block 内线程共享的片上内存 L2 cache GPU 全局共享缓存，连接 SM 和 HBM HBM / global memory 显存，容量大但访问慢 关键事实：Tensor Core 上的矩阵乘法吞吐远高于普通 CUDA core 上的一般浮点操作。因此高性能实现会努力把计算组织成 Tensor Core 友好的矩阵乘法形式。\nCUDA 执行模型：thread、block、warp # 三层执行抽象：\nthread: 最小逻辑执行单元 block: 一组 threads，可以共享 shared memory grid: 一个 kernel launch 中的所有 blocks 硬件调度中还有 warp：\nwarp = GPU 实际调度执行的线程组，通常包含 32 个连续 threads GPU 使用 SIMT 模型（Single Instruction, Multiple Threads）：同一个 warp 中的线程通常执行同一条指令，但操作不同数据。\n这解释了为什么分支会拖慢 GPU——如果同一个 warp 中一部分线程走 path A，另一部分走 path B，GPU 会 serially 执行两个路径：先执行 path A 并 mask 掉走 path B 的线程，再执行 path B 并 mask 掉走 path A 的线程。总时间 ≈ 两个分支时间之和。这叫 branch divergence。\n关键细节：diverge 时线程不是 idle，而是被 mask 掉不执行当前指令。所以同一个 warp 内应尽量避免数据相关的分支。\nGPU 适合规则、密集、统一的计算模式；不适合大量不规则分支。 GPU 内存层次 # 从快到慢：\nregisters ↓ shared memory / L1 cache ↓ L2 cache ↓ HBM / global memory 层级 位置 特点 register thread 私有 最快，但容量最小 shared memory block 内共享 很快，适合 tile 复用 L1 cache SM 附近 缓存局部访问 L2 cache GPU 全局 所有 SM 共享 HBM/global memory 显存 容量大，带宽高但延迟远高于片上内存 重要原则：一次从 HBM 读入的数据，最好在 register/shared memory 中被尽可能多次复用。\n如果一个算子反复从 HBM 读写中间结果，即使 FLOPs 不多，也可能非常慢。这也是 FlashAttention 的核心动机：减少 attention matrix 在 HBM 中的读写。\nCompute scaling 快于 memory scaling # 现代 GPU 计算能力增长非常快（尤其 Tensor Core），但显存带宽和内存访问速度没有以同样速度增长。\n很多程序不是算不动，而是数据搬不够快。 类型 含义 compute-bound 主要受计算单元峰值 FLOPs 限制 memory-bound 主要受内存带宽或数据搬运限制 如果一个 kernel 每读入很多 bytes 只做很少 FLOPs，它很可能是 memory-bound。反过来，如果每读入一份数据能做大量计算并反复复用，就更可能接近 compute-bound。\nArithmetic intensity 与 roofline model # Arithmetic intensity 衡量每搬运一个 byte 数据能做多少计算：\nArithmetic Intensity = FLOPs / Bytes moved 可达到的吞吐可以用 roofline model 估计：\nAttainable FLOP/s = min(Peak FLOP/s, Memory Bandwidth × Arithmetic Intensity) 低 arithmetic intensity：受内存带宽限制，memory-bound 高 arithmetic intensity：受峰值计算限制，compute-bound 转折点称为 ridge point：\nridge point = Peak FLOP/s / Memory Bandwidth 如果某个算子的 arithmetic intensity 低于 ridge point，提高计算单元数量不一定有帮助；更应该减少数据搬运、做 operator fusion、tiling 或 recomputation。\n数值直觉：A100 的 FP32 peak 约 19.5 TFLOP/s，HBM bandwidth 约 2 TB/s，ridge point ≈ 10 FLOPs/byte。BF16 Tensor Core peak 约 312 TFLOP/s，ridge point ≈ 156 FLOPs/byte。所以低精度 matmul 更容易 compute-bound，而 elementwise/softmax 更容易 memory-bound。\nMatmul 为什么容易快 # 矩阵乘法 $C = AB$（$A,B,C \\in \\mathbb{R}^{N \\times N}$）：FLOPs ≈ $2N^3$，理想读写数据规模约为 $O(N^2)$，所以理想 arithmetic intensity 约为 $O(N)$。当 N 足够大时，矩阵乘法可以有很高的数据复用，容易接近 compute-bound。\n但 naive matmul 不一定快：如果每次计算都从 global memory 重复读取 A 和 B 的元素，实际 bytes moved 会大幅增加。解决方法是 tiling：\n1. 把 A 和 B 切成 tile 2. 把 tile 加载到 shared memory / registers 3. 在片上反复复用 tile 4. 计算出 C 的一个 tile 让 GPU workload 变快的六类技巧 # 避免 control divergence：让同一个 warp 中的线程执行尽可能一致的控制流，避免数据相关的复杂分支。\n使用低精度计算：FP16/BF16/FP8/INT8 减少内存读写、提高缓存有效容量、利用 Tensor Cores。但会带来数值稳定性问题，需要 loss scaling、FP32 master weights、混合精度策略。\nOperator fusion：未融合的 elementwise chain 每步都读写 HBM；fusion 后中间结果留在 register 中，大幅减少 HBM traffic。\nRecomputation：用额外计算换更少存储——前向不保存某些中间激活，反向时重新计算。核心权衡：多做一点 FLOPs，少占很多 HBM。memory-bound 时这个交换往往值得。\nCoalesced memory access：相邻线程访问连续内存，一个 warp 的访问可以合并成较少的 memory transactions。\nTiling：把一小块数据搬到快内存中，并在写回 HBM 前尽可能多次复用。在 matmul 中复用 A/B 子矩阵；在 attention 中分块计算 QK^T 和 softmax，避免完整 attention matrix 落到 HBM。\nOccupancy 的直觉 # Occupancy 指一个 SM 上实际活跃 warp 数量相对于理论最大数量的比例。高 occupancy 的意义：当某些 warp 等待内存时，SM 可以切换到其他 ready warp 继续执行，从而隐藏延迟。\n但 occupancy 不是越高越好。影响 occupancy 的资源：每个 thread 使用的 registers 数量、每个 block 使用的 shared memory、block size、硬件最大 active warps/blocks 限制。如果为了提高 occupancy 而牺牲数据复用，可能反而变慢。\nAttention 为什么容易 memory-bound # 标准 attention：$\\text{Attention}(Q, K, V) = \\text{softmax}(QK^T / \\sqrt{d}) V$。如果序列长度为 n，attention score matrix 大小为 $n \\times n$。\n普通实现：\n1. 计算 S = QK^T 2. 把 S 写入 HBM 3. 从 HBM 读出 S，做 softmax 4. 把 P = softmax(S) 写入 HBM 5. 从 HBM 读出 P 和 V，计算 O = P V 这会产生大量 HBM traffic，尤其当 n 很大时，中间矩阵 S 和 P 都是 $O(n^2)$。问题不只是 FLOPs，而是中间 attention matrix 太大、频繁读写 HBM。此外 softmax 包含跨 key 维度的 reduction（求 max、求 sum），需要 warp 内线程通信，进一步增加延迟。\nFlashAttention 的核心思想 # FlashAttention 是 IO-aware exact attention。它不改变 attention 的数学定义，而是改变计算顺序：\n把 Q, K, V 分块 在 SRAM/shared memory/registers 中计算局部 QK^T 用 online softmax 维护正确的归一化统计量 逐块累积输出 O 避免把完整 S 或 P 写入 HBM 关键词：tiling、online softmax、recomputation、HBM traffic reduction、exact attention。\n重要性：它不是近似 attention，而是通过系统优化得到数学等价的结果。\nOnline softmax 回顾 # 标准 stable softmax 对一行 score：\nm = max_i s_i l = sum_i exp(s_i - m) p_i = exp(s_i - m) / l 分块 attention 中不能一次性拿到整行 score，需要 online softmax。假设旧状态为 $m_{old}, l_{old}$，新 block 的最大值和分母为 $m_B = \\max(S_B)$，$l_B = \\sum \\exp(S_B - m_B)$，合并：\n$$m_{new} = \\max(m_{old}, m_B)$$$$l_{new} = l_{old} \\cdot \\exp(m_{old} - m_{new}) + l_B \\cdot \\exp(m_B - m_{new})$$如果还维护未归一化输出 $a$：\n$$a_{new} = a_{old} \\cdot \\exp(m_{old} - m_{new}) + a_B \\cdot \\exp(m_B - m_{new})$$$$\\text{output} = a_{new} / l_{new}$$当新的最大值出现时，旧结果不是被忽略，而是整体乘以 $\\exp(m_{old} - m_{new})$——等价于把所有旧项从减去旧最大值重新换成减去新最大值。\n为什么 FlashAttention 能减少内存 # 普通 attention 的瓶颈：需要显式存储 $S = QK^T$ 和 $P = \\text{softmax}(S)$，二者都是 $O(n^2)$。\nFlashAttention 的策略：不把完整 S/P 写入 HBM，只保存每行的 $m$、$l$ 和最终 $O$，每个 block 内临时 score 用完即丢弃。反向传播时根据保存的少量统计量重算局部 score 和 softmax（recomputation）。\n核心 tradeoff：多做一些计算，显著减少 HBM 读写和中间存储。在 attention 这种 memory-bound 场景中，这个 tradeoff 很划算。\n如何用这些直觉分析 kernel # 做 profiling 或写 kernel 时，应该反复问：\n这个 kernel 是 compute-bound 还是 memory-bound？ 数据是否 coalesced 访问？ 有没有多次读写 HBM 的中间结果？ 是否可以 fusion？ 是否可以 tiling 到 shared memory/registers？ 是否可以用 recomputation 换内存？ Tensor Core 是否被充分利用？ occupancy 低是瓶颈还是合理的资源权衡？ 几个典型判断：\nforward pass 中 GEMM kernel 占比高，因为 matmul 的 arithmetic intensity 高，容易 compute-bound。 完整 training step 中 elementwise/reduction/optimizer kernel 显著增加，因为 backward 需要逐元素梯度计算，AdamW 需要逐元素更新一阶/二阶动量，这些操作 arithmetic intensity 低，更偏 memory-bound。 attention_softmax 的 runtime 可能比 attention_scores_matmul 更长，因为 softmax 受 memory bandwidth 和 reduction 通信限制，而 matmul 被 cuBLAS 高度优化后 compute throughput 更高。 mixed precision 中 GEMM 加速最多（Tensor Core），elementwise/reduction 可能变化不大（带宽限制 + 数值稳定性要求）。 GPU 优化本质上是围绕硬件约束重新组织数学计算——这比单纯看代码更重要。\n术语表 # 术语 中文理解 GPU 面向高吞吐并行计算的处理器 TPU 面向矩阵/张量计算的专用加速器 SM GPU 中执行 thread block 的主要计算单元 Warp GPU 调度单位，通常 32 个 threads SIMT 单指令多线程执行模型 Register 线程私有最快存储 Shared memory block 内共享片上内存 HBM 高带宽显存，也常被称为 global memory Tensor Core 专门加速矩阵乘法的硬件 Occupancy SM 上活跃 warp 占理论上限的比例 Arithmetic intensity FLOPs / bytes moved Roofline model 用计算峰值和内存带宽估计性能上限的模型 Memory-bound 性能主要受数据搬运限制 Compute-bound 性能主要受计算峰值限制 Operator fusion 合并多个算子减少中间读写 Recomputation 用额外计算换更少存储 Coalescing 相邻线程访问连续内存以合并 transaction Tiling 分块计算并在快内存中复用数据 FlashAttention IO-aware exact attention 实现 一句话总结 # LLM 系统性能不只由 FLOPs 决定，更由数据如何在 GPU 内存层次中移动决定。 GPU 优化的关键是：让大量线程规则并行执行，让 Tensor Cores 吃满矩阵乘法，让数据尽量留在 register/shared memory 中复用，减少 HBM 往返。FlashAttention 是这一思想的代表——用 tiling、online softmax 和 recomputation 重新组织 attention 计算，在保持数学等价的同时显著降低 memory I/O。\n","date":"2026年6月13日","externalUrl":null,"permalink":"/notes/cuda/cs336-gpu-attention/","section":"笔记","summary":"从模型结构转入系统视角：为什么 LLM 必须依赖 GPU、GPU 的执行模型与内存层次、arithmetic intensity 与 roofline model，以及如何用这套性能模型解释 FlashAttention 这类 IO-aware 算法。","title":"GPU 架构与 FlashAttention：CS336 Lecture 5 笔记","type":"notes"},{"content":"","externalUrl":null,"permalink":"/authors/","section":"Authors","summary":"","title":"Authors","type":"authors"},{"content":"","externalUrl":null,"permalink":"/categories/","section":"Categories","summary":"","title":"Categories","type":"categories"},{"content":"","externalUrl":null,"permalink":"/series/","section":"Series","summary":"","title":"Series","type":"series"},{"content":" ","externalUrl":null,"permalink":"/","section":"首页","summary":"","title":"首页","type":"page"}]