[{"content":"这个系列是什么 这是一份学习笔记 —— 我自己在搞清楚现代 LLM 到底是怎么被 serve 起来的，主要靠跟 Claude 聊，然后把真正想明白的部分写下来。文章本身是用一种\u0026quot;发现之旅\u0026quot;的、比较自信的口吻写的，但底下其实就是一个普通人在公开学习。\n下面这张表是活的 —— 文章会随着发布翻状态，每次讨论挖出新坑就会往里加。\n文章列表 # 标题 状态 链接 01 LLM 从头到尾走一遍 —— 鸟瞰整张网络、单个 block 的内部、生成一段文字的完整过程，以及由此自然冒出的、贯穿整个系列的问题 [next] — 02 Tensor parallelism 心智模型：从零搭起 [done] 阅读 → 03 在一个 transformer block 中完整走完一遍 Tensor Parallelism —— 先全用 column-parallel 看通信怎么爆掉，再配上 row-parallel 落到每个 block 两次 all-reduce [done] 阅读 → 04 一次 forward 怎么塞下很多个 request —— varlen attention，只考虑 prefill，TP 完全没动到 [done] 阅读 → 05 ORCA 和 chunked prefill —— iteration-level 调度先把进出 batch 的边界问题解决掉；chunked prefill 再给每次 iteration 的开销封顶，免得一个长 prompt 把整台引擎卡住 [done] (EN) read → 06 Prefill 和 decode 拆机 —— 两个阶段在 roofline 的两侧，承认这种 asymmetry 之后，共用一池 GPU 就不再是折中，而是在跟公式硬扛 [done] 阅读 → 07 拆机后的工程问题 —— KV cache 跨 fabric 怎么传（NVLink、NVSwitch、IB、PCIe）、tiered memory pool（HBM、DRAM、SSD）、跟 prefill 怎么 overlap、按拓扑路由 [next] — 08 Pipeline parallelism —— 切 跨 block 而不是 block 内部，以及它带来的 bubble；prefill pool 为什么想要它 [planned] — 09 MoE 和 expert parallelism —— FFN 变成 routed 之后改了什么 [planned] — 10 PagedAttention —— 把 KV cache 当虚拟内存做，分块而不是连续，跨 request 还能 copy-on-write [planned] — 11 Sequence 和 context parallelism —— 把一个 request 切到多张 GPU 上，ring attention，长上下文那一招 [planned] — 12 FlashAttention —— tiled online softmax，为什么 [L × L] 的 score matrix 根本不需要存在 [speculative] — 13 FlashDecoding —— bandwidth 压力下让 1 × L_kv 的 decode-attention 跑快 [speculative] — 14 GQA 和 MLA —— 更少的 KV head、更小的 KV cache、更快的 decode（以及模型代价） [speculative] — 15 Speculative decoding —— 一个小模型出 draft，大模型来 verify，两次 forward 价钱办一次的事 [speculative] — 16 KV 压缩 —— 量化、eviction，能丢的和不能丢的 [speculative] — 状态说明 [done] 已发布、有链接 · [done] (EN) 英文版已发布，中文版还在路上 · [next] 正在写 · [planned] 排队中，肯定会写到 · [speculative] 一个值得挖的坑 —— 可能填、可能不填，但问题本身有意思\n几条贯穿多篇的线索 下面这些观察会在多篇文章里反复出现，读的时候可以放在脑子里：\nTP 出乎意料地不挡道。 Request batching 没动到它（Article 04），continuous batching + chunked prefill 也没动到（Article 05）。但 PP 和 MoE 会 跟 TP 有一些有意思的互动 —— 这也是为什么这两篇在地图上排得比较靠前。 KV cache 是文章 05 之后的连结组织。 它一旦在 decode 里登场就不会离开了；它也是长上下文之所以困难的根源。 Decode 把瓶颈翻了个面。 02–04 都假设 prefill，那时候 compute 是主要约束。一旦 decode 进来（Article 05 起），weight 读取的 bandwidth 变成卡脖子的那个 —— 这也是后面几乎每个优化（FlashDecoding、GQA、prefill/decode 拆机、speculative decoding）的动机。 Modeler 的选择反复在 serving 这边承重。 Multi-head 的独立性让 TP 通信免费，也让 request batching 通信免费，到 GQA/MLA 那篇还会再出现一次。这是一条值得留意的反复出现的主题。 ","permalink":"https://wgzesg.github.io/llm_stories/zh/posts/00-roadmap/","summary":"这个系列是什么，以及一份还在生长的文章地图 —— 已发布、在写、留给未来的自己去填的坑都列在这。","title":"Roadmap：这个系列要写些什么"},{"content":"你看到 LLM 写出来的每一段回复，都是一个 token 一个 token 地生成出来的 —— 同一个固定大小的模型，反反复复地在自己刚吐出来的 output 上跑了一遍又一遍。不是只有那段最精彩的，而是每一段都这么来。这个模型本身叫 transformer，外面那个反复调用它的循环则基本上是机械式的记账工作。把这个模型搞懂、把那个循环搞懂，这件事就吃透了。系列后面的所有事 —— 把模型切到多 GPU 上、让很多用户共享一次 forward、让超长 prompt 装得下、让生成更快 —— 都建立在这两块之上。\n这一篇从三个 zoom level 把模型打开：\n整个模型，从头到尾 —— 进什么、出什么、中间发生了什么。 拉近一层 —— 那个标着 \u0026ldquo;transformer block\u0026rdquo; 的东西，里面究竟是什么。 完整循环 —— 这个一次只生成一个 token 的模型，怎么被用来吐出一长段回复的。 为了讨论尽量通用，我们一直用符号（d、L、h 之类）而不是具体数字 —— 因为各家模型的结构相通，但具体数字千差万别。不同模型尺寸不一样，但骨架都长这样。具体数字留给后面的文章 —— 等到它们真的承重的时候再说。\n边读边会冒出一些自然的问题 —— 比如*\u0026ldquo;等等，每生成一个 token 模型就要把前面那么多事重新做一遍？\u0026quot;或者\u0026ldquo;那如果模型大到一张 GPU 装不下怎么办？\u0026quot;*这些问题，正是后面整个系列要逐一拆开来回答的。每个问题最后都会有自己的一篇。\nPart I —— 整个模型，从头到尾 1. token 进来，下一个 token 出去 塞给模型一段话 —— 比如 \u0026quot;the quick brown fox jumps over\u0026quot; —— 让它接着往下写。它实际上到底做了哪些事？从头到尾六步。\n1. Tokenize。 第一步把字符串切成一小段一小段，每一段叫一个 token。每个 token 都是一个小整数 ID —— 因为模型底层只会做算术，处理不了\u0026quot;字\u0026quot;本身。粗略地说：常见的短词通常一个 token，生僻词或长词会被拆成几段。token 的数量我们记作 N。\n2. Embed。 每个整数 ID 拿去查一张巨大的表 —— embedding table。表里每一行对应词表里的一个 token，每一行是 d 个数字组成的向量（d 是模型自己挑的一个超参，叫 hidden dimension。真实模型里 d 一般在几千这个量级）。N 个 token 查完之后，原本一串 N 个整数 ID 变成了一个形状 [N × d] 的 tensor：N 行，每行 d 个数。\n为什么要换成向量、不直接用整数 ID？因为模型底层只会做线性代数，整数 ID 之间没有任何有用的几何关系 —— token 5 不会因为是连续整数就比 token 100 离 token 6 \u0026ldquo;更近\u0026rdquo;。embedding table 给每个 token 在 d 维空间里安排了一个学到的位置：意思相近的 token 落在附近，没什么关系的 token 离得远。每一行可以理解成模型对那个 token 的\u0026quot;第一印象\u0026rdquo; —— 还没看到它在句子里的上下文之前，凭空对它的感觉。\n（词表大小记作 vocab，一般几万这个量级。所以 embedding table 自己就是一个 [vocab × d] 的矩阵 —— 这本身就是不小一坨参数，§2 里再聊。）\n3. 一摞 transformer block。 这个 [N × d] 的 tensor 接着穿过 L 个 transformer block，一个摞一个。每个 block 都会读一遍完整序列，把不同 position 之间的信息混合一下，再写回一个更精炼的版本。关键是：每个 block 的 input 和 output 形状完全一样，都是 [N × d] —— 只是行里的内容被改了。\nL 个 block 都过完之后，每一行已经离最初那个起点很远了。它代表的不再是这个 token 的脱离上下文的、通用含义，而是它在这个具体序列里的含义。block 为什么能这样一直摞下去，§2 专门聊；§Part II 会把一个 block 拆开来看。\n4. Final norm。 一摞 block 顶上还有一个小小的归一化步骤 —— 算是个收尾的整理。形状不变，进什么形状出什么形状。\n5. LM head。 一个 linear layer 把每一行从 d 维投回去，每一行变成 vocab 个数 —— 词表里每个 token 一个数。output 形状 [N × vocab]。每一行是一长条对整个词表的\u0026quot;打分\u0026rdquo;。这种原始分数叫 logits。位置 i 上 token t 的 logit，是模型对*\u0026ldquo;在 position i 上下一个 token 是 t 有多合理\u0026rdquo;*那种原始的、没归一化过的回答。\n6. Softmax → sample。 我们真正要的是最后一行 —— 也就是最后一个输入 token 之后那个位置，那里放着模型对\u0026quot;下一个该是什么\u0026quot;的预测。Softmax 把那一行的 logits 拧成一个干净的概率分布 —— 全是正数，加起来等于 1。从这个分布里采样一个 token，这就是模型对下一个 token 的猜测。\n整个 stack 拍下来：\nthe model, end to end token IDs (integers) shape: [N] embedding lookup [vocab × d] [N × d] L transformer blocks [N × d] in, [N × d] out, repeated block 1 block 2 ⋮ ⋮ block L−1 block L [N × d] final LayerNorm LM head [d × vocab] [N × vocab] logits softmax(last row) → next-token distribution 所以整个模型本质上就是一个函数：吃 N 个 token，吐回一个对第 N+1 个 token 该是什么的概率分布。其它所有东西 —— 那些聊天式的回复、长篇回答、聊天 UI 里一个字一个字蹦出来的流式输出 —— 都是把这个函数反复调用得来的。这个循环，Part III 来讲。\n2. block 为什么能一直摞下去：stream-processor 模式 一句话讲清楚 transformer：一摞 L 个完全同款的 \u0026ldquo;stream processor\u0026rdquo;，吃一个固定形状的 token 流，加工一下，往下传。这个形状就是 [N × d]。同样的形状进、同样的形状出，重复 L 次。\n这条性质为什么重要？两个理由，后面整个系列都在反复用：\n它让模型可以靠\u0026quot;摞\u0026quot;来变大。 想要一个更大的模型？多摞几个 block 就行。一个小型开源模型和一个巨大的旗舰模型，从这个 zoom level 看几乎是一模一样的 —— 同样的六步 pipeline、同样的 block 结构，只是 L 不同（d 也稍微宽一点）。同一份菜谱，放大版本。 它让所有下游工具都不用关心\u0026quot;在第几层\u0026quot;。 一个 block 根本不知道自己是第 1 个还是第 32 个，所以任何接触 block 的工具（切到多 GPU 的 splitter、做 batching 的 batcher、做 scheduling 的 scheduler）也都不用关心。整摞 block 是一片整齐的 substrate，工具直接在上面操作。 （这个想法在别的地方你可能也见过 —— Unix pipe、音频插件、图像处理流水线。同样的形状进、同样的形状出，想摞多少摞多少。）\n把形状钉得更死一点：回看 §1 的六步，过了 tokenization 之后，中间每一步进出的都是同一个 [N × d] tensor。\nembedding 把 N 个 token ID 变成一个 [N × d] tensor。 每个 transformer block 读 [N × d]、返回 [N × d]。 final norm 读 [N × d]、返回 [N × d]。 只有最顶上的 LM head 改了宽度 —— 把它换回 vocab 那么宽。 管道中段，形状从来不变。变的是内容 —— 每个 block 都在精炼这些行，慢慢搭出越来越丰富、越来越能感知上下文的表示 —— 但从最底层到最顶层，几何结构始终是 [N × d]。\n后面会反复用到的几个符号：\nN —— 当前这段序列的长度。每次请求都不一样 —— 它是 input 的属性，不是模型的属性。 d —— hidden dimension。流过 stack 的 tensor 每一行的宽度。 L —— 摞了多少个 transformer block。 vocab —— 模型认识多少种不同的 token。决定了 embedding table 和 LM head 的宽度。 Part II 还会再见到两个：h（一个 block 里的 attention head 数量）和 d_head（每个 head 多宽）。\nPart II —— 把一个 block 打开 3. 一个 block 拍平来看 现在我们打开 L 个 transformer block 里的一个。好消息是：它们内部结构都一样 —— 不同 block 学到的参数不同，但接线方式一模一样。看懂一个，就看懂了所有 L 个。\n一个 block 分成两半，每一半都被一个 residual connection 包起来（每一半底下那个小小的 + —— 一会儿就解释）：\none block — two halves, each wrapped in a residual input [N × d] residual attention sub-layer LayerNorm 1 tidy-up QKV projection d → 3d, split into Q, K, V multi-head attention mixes across positions output projection d → d + [N × d] residual FFN sub-layer LayerNorm 2 tidy-up FFN-up d → 4d activation (GeLU) pointwise nonlinearity FFN-down 4d → d + output [N × d] 这两半就是 block 干的两件主要的事：一个 attention sub-layer，一个 FFN（feed-forward network）sub-layer。其它那些零件（LayerNorm、activation、+）是小一些的胶水。\n每个零件大致在干什么：\nLayerNorm 是一个归一化步骤 —— 对 tensor 的每一行，把里面的数重新缩放，让它们的均值和方差落在一个干净的范围里。便宜、按行做、纯粹是为了在数字穿过很多层之后不让它们漂到奇怪的数量级去。可以当成一个\u0026quot;整理\u0026quot;步骤。 residual + 的意思是：把进入这一半之前的东西和这一半算出来的东西加在一起。所以每一半算出来的其实是一个 delta —— 在已有表示上做一次精炼，而不是整个换掉。这就是我们能摞很多个 block 而信号一路不糊的原因。 QKV projection 就是三个 linear layer 合并成一次大 matmul。它给 input 套三个不同的 weight matrix，产出三个 tensor —— Q（queries）、K（keys）、V（values）—— 每个形状都是 [N × d]。 Multi-head attention 是整个模型里唯一让信息在 token 之间流动的步骤。它是 block 的主角 —— §4 讲它真正在算什么，§5 讲为什么前面要加 \u0026ldquo;multi-head\u0026rdquo;。 output projection 是最后一个 linear layer，把 attention 的输出整合成能被 residual + 直接吸收的形态。 FFN-up 和 FFN-down 是中间夹一道非线性的两个 linear layer。它俩合作把每个 token 的 d 维表示先扩到 4d、过一道按元素的非线性、再压回 d。不在 token 之间混 —— 每个 token 各自处理自己。 同样的形状进、同样的形状出 —— §2 那条口诀。摞很多个，就是模型的主体。\n4. attention 到底在算什么 \u0026ldquo;attention 在 position 之间混合\u0026quot;这句话我们说过好几遍了，但一直没讲怎么混合。这一节补上。\n对每个 position，模型从这个位置 [d] 维的那一行里生出三个向量：\n一个 query Q —— \u0026ldquo;我在找什么？\u0026rdquo; 一个 key K —— \u0026ldquo;我能提供什么？\u0026rdquo; 一个 value V —— \u0026ldquo;如果你决定关注我，这是我想传过去的实际内容\u0026rdquo; （这正是 QKV projection 在做的事 —— 三个 linear layer，每个负责 Q、K、V 之一，融合成一次 matmul。）\n要更新 position i 那一行，模型做三件事：\n算 score。 拿 i 的 query 跟每一个 position 的 key 做点积。点积大 = 两个向量方向接近 = \u0026ldquo;这个 position 对 i 来说有意思\u0026rdquo;。点积小（或者负数）= 不感兴趣。最后得到 N 个分数 —— 每个 position 一个。 score 变成权重。 把这些分数过一道 softmax，得到 attention weight —— 全是正数、加起来等于 1。位置 j 上权重高，意思就是*\u0026ldquo;i 很关心 j\u0026rdquo;；权重低，就是\u0026ldquo;i 基本上忽略 j\u0026rdquo;*。 对 value 做加权平均。 拿这些权重，对每个 position 的 value 向量做加权求和。这个和，就是 i 这一行更新后的表示。 一句话：position i 的新一行，是所有 position 的 value 向量的加权平均；权重由 i 的 query 跟每个 position 的 key 的匹配度决定。\n就这 —— 这就是 attention 全部的机械内容。block 里其它一切（LayerNorm、FFN、residual）都是为了支撑这一件事存在的基础设施。它也是整个模型里唯一让信息在 token 之间流动的步骤。把 attention 拿掉，模型就分不清 \u0026ldquo;fox\u0026rdquo; 和 \u0026ldquo;the\u0026rdquo; 在不在同一个句子里了。\n这套机制后面还要再补两个细节：\n§5 —— Heads。 attention 不是在 d 维的整段 feature 上跑一次的，而是在不同的 feature 切片上并行跑多次。 §6 —— Causal mask。 position i 其实不能 attend 到所有 position，只能看 j ≤ i。为什么要这样，§6 讲。 FFN sub-layer 相比之下简单很多：每一行都过同样的两个 linear layer 加一道非线性，跟其它行互不相干。FFN 不在 position 之间混 —— 那是 attention sub-layer 的活。\n所以每个 transformer block 的节奏就是：position 之间混（attention），然后 feature 之间混（FFN）。 重复 L 次。\n5. Heads 关于 attention 还有一件小事，但后面会很重要：它不是在 d 维整段 feature 上跑一次，而是切成 h 段并行跑 h 次。\nQKV projection 生出形状都是 [N × d] 的 Q、K、V 之后，我们沿 feature 维把每个reshape 成 h 组，每组宽度 d_head = d / h。每一组就是一个 head。每个 head 在自己那一片 feature 上跑一遍 §4 的 attention —— 自己的 query、自己的 key、自己的 value。所有 head 的输出再 concat 回 [N × d]，喂给 output projection。\nmulti-head attention: reshape, per-head, concat Q, K, V [N × d] reshape h heads, each d_head wide [N × h × d_head] attention (§4) h attention outputs [N × h × d_head] concat to output proj [N × d] each head runs §4\u0026rsquo;s attention algorithm on its own slice — independently of the others 真实模型里 h 和 d_head 一定凑出来正好等于 d —— 一般是几十个 head，每个一百多宽。\n模型设计角度看：不同的 head 可以学着去关注不同种类的东西。有些 head 最后在追踪短距离的句法关系（\u0026ldquo;这个代词到底指代哪个词？\u0026quot;），有些跟踪更远距离的模式。多个 head = 多个看 \u0026ldquo;应该看哪里\u0026rdquo; 的视角。\n系统层面就更直白：head 之间是独立的。 Head 0 在 attention 里不和 Head 1 说话。每个 head 在自己那块 feature 上各算各的，各出各的输出。\n这种独立性只是模型设计的一条性质而已 —— 但它对后面所有事情都是承重墙。Article 02 直接利用这条性质，把整个模型一刀切到两张 GPU 上：一半的 head 在一张卡，另一半的 head 在另一张卡，attention 期间它们之间根本不需要通信。\u0026ldquo;模型大到一张 GPU 装不下怎么办\u0026quot;的整个故事，起点就在这里。\n6. Causal mask attention 里还有一条规则没讲，但它是必不可少的：position i 在 attend 的时候，只能看 j ≤ i 的位置。j \u0026gt; i 的位置会被 mask 掉 —— 它们的 attention score 在过 softmax 之前会被强行置成 −∞，过完 softmax 权重就变成 0，对 i 的 output 没有任何贡献。\n为什么要有这条规则，原因来自训练。模型是按\u0026quot;预测下一个 token\u0026quot;一条一条训练的：喂一段序列进去，让模型从每个 token 之前的所有东西去预测它的下一个 token。如果 position i 在 attention 里能偷看 position i+1，那它就等于可以作弊，直接读答案。mask 就是用来强制\u0026quot;不许往后看\u0026quot;的。\nmask 还有两个后续影响值得专门点名，后面都会用到。\n第一，它让 Part III 里那个生成循环成立：position N+1 的 token 只依赖于 token 1..N，反过来不会。所以我们可以按顺序一个一个生成新 token，从来不需要回头修改一个已经算好的 token。这条性质让\u0026quot;一个一个 token 地生成长回复\u0026quot;这件事根本能成立。\n第二 —— 也是更大的那个 —— mask 意味着旧 token 的活永远不需要重做。position 5 的 hidden state，无论整段序列是 5 个 token 长还是 500 个 token 长，都是同一个；后面任何新 token 都伸不回来把它改了。这种\u0026quot;算完就不再变\u0026quot;的性质，是我们能不能想到 \u0026ldquo;把之前算过的存下来下次复用，而不是每次 forward 都从头算\u0026rdquo; 的前提。没有 mask，每来一个新 token 就要把前面所有东西重新过一遍。有了 mask，我们才能想着按顺序处理 token，把已经算过的记住就好 —— 这正是 §10 会落到的问题，也是后面整个系列里最承重的一类优化之一。\n7. 一张图把整个 block 装进去 到这里我们已经把一个 transformer block 的所有零件都打开过了 —— 两个一半（§3）、attention 的 Q/K/V 机制（§4）、按 head 切（§5）、causal mask（§6）。下面这张图把它们按一次完整的执行画出来，每一步都标了 tensor 形状。\n先扫一遍感受整体流向，之后系列里看到*\u0026ldquo;那个 [h × N × N] 的 score matrix\u0026rdquo;或者\u0026ldquo;按 head 切开的 reshape\u0026rdquo;*时，回来对一下 —— 你要在脑子里想象的就是这张图。\ninside one block — every operation, every shape input [N × d] residual LayerNorm 1 [N × d] QKV projection Q K V each [N × d] reshape Q, K, V along feature dim into h heads each [N × h × d_head] multi-head attention (per head, in parallel) Q · Kᵀ / √d_head [h × N × N] scores + causal mask (future → −∞) [h × N × N] softmax (along last dim) [h × N × N] weights weights · V [N × h × d_head] concat heads back into [N × d] [N × d] output projection + [N × d] residual LayerNorm 2 [N × d] FFN-up (d → 4d) [N × 4d] activation (GeLU) [N × 4d] FFN-down (4d → d) + output [N × d] 有三点值得停下来看一眼：\n形状从 [N × d] 进，从 [N × d] 出 —— §2 那条口诀。在一个 block 内部，tensor 会短暂地变成别的形状（FFN 中间是 [N × 4d]，attention score 是 [h × N × N]） —— 但这些都是瞬态的。block 最后总会回到 [N × d]，这样下一个 block 才能接得上。 [h × N × N] 这个 score matrix 是会让人吃惊的那个。 它的大小按序列长度的平方长。N 小的时候没事，N 一大就难处理 —— 长序列的代价最后就栽在这里。现在留意一下，后面的文章会回来收拾它。 每个 residual + 把那一半的 input 重新加回 output。 所以每一半算的其实是 delta —— 一次微调，而不是把整段表示整个换掉。这就是为什么我们能摞很多 block 而信号不崩。 Part III —— 用模型来生成 8. 一次 forward 给你一个 token §1 那个模型，吃 N 个 token，吐回一个对下一个 token 的概率分布。一个 token。不是一整句，连半句也不是 —— 就一个对下一个 token 的猜测。\n但我们已经习惯 LLM 一段一段地回复。一次只能生成一个 token 的模型，怎么吐出一整段？跟你猜的一样：反复跑，把自己的 output 当成下一次的 input 喂回去。\n具体来说：\n起点：prompt —— 一段长度为 N 的序列。 跑一次 forward。得到下一个 token（也就是位置 N+1）应该是什么的分布。 从这个分布里采样（或者直接挑最高概率那个，\u0026ldquo;argmax\u0026rdquo;）。位置 N+1 上就有了一个 token。 把它追加到序列后面。序列长度变成 N+1。 在完整 的 N+1 长序列上再跑一次 forward。得到位置 N+2 的分布。 采样、追加。序列长度变成 N+2。 重复，直到模型采样到一个特殊的 end-of-sequence token（训练时模型就被教会：认为回答结束时，发出这个 token），或者撞到你设定的长度上限。 循环长这样：\ngeneration loop: sample, append, repeat prompt length N forward on length N sample token N+1 from last-row softmax sequence length N+1 feed the appended sequence back in sequence length N+1 forward on length N+1 sample token N+2 sequence length N+2 ⋮\nuntil model emits end-of-sequence or until a length cap is hit 数学上讲，整个生成流程就是这样。你用过的任何基于 LLM 的系统，吐出来的每一个 token，都是从一个长这样的循环里出来的。\n9. 第一个让人不舒服的观察 我们走一遍：从一个长度为 N 的 prompt 开始，生成 K 个新 token，总成本是多少。\nForward 1 跑在 prompt 上：长度 N。 Forward 2 跑在 prompt + 1 个新 token 上：长度 N+1。 Forward 3：长度 N+2。 … Forward K：长度 N+K−1。 每一次 forward 几乎都把上一次做过的事再做一遍。Forward 2 input 的前 N 个 token 和 Forward 1 input 一模一样 —— 但模型还是从头在每一个 position 上跑了一遍每个 block，就像它从没见过这些 token 一样。\n总的算下来，工作量大致按 (N + K)² / 2 涨 —— 总序列长度的平方。其中绝大部分都是在重新计算根本没变的东西。在序列末尾加一个新 token，并不会改变前面任何 token 的表示。前面那些 token 还是原来那个 prompt，加上这次之前已经采样出来的那几个 token 而已。它们身上没有任何东西需要重新算。\n于是一个很显然的问题就挂在那里：这些重算真的有必要吗？ 显然没有。但是不重算也不是白来的 —— 这意味着我们要在两次 forward 之间存某种中间状态。然后这又冒出一串新问题：到底要存什么状态？放哪？它有多大？随着对话变长它怎么涨？\n这种问题，正是这个系列后面要拆开来研究的。\n10. 一张问题地图 后面这个系列其实就在追两条主线，关于\u0026quot;怎么运行一个 LLM\u0026quot;的绝大多数实际问题，都能归到其中一条。\n主线 1 —— 让一次 forward 装得下。 一次穿过这摞 block 的 forward，可能在好几条尺度上都\u0026quot;太大\u0026quot;了：weight 装不下一张 GPU、算一次时间太久、attention 内部太吃内存。这条主线下的文章，主题都是把工作在空间上切开，让一次 forward 能落在你手里这套硬件上跑下来。\n模型本身可以很大。 摞够多 block（L 大）、d 又够宽，光是 weight 就装不下一张 GPU。怎么把一次 forward 切到多张 GPU 上？(Article 02 和 03 —— 用的正是我们在 §5 搭起来的 head 独立性。) prompt 本身可以很大。 §7 那个 [h × N × N] 的 score matrix 按序列长度的平方涨。prompt 一长，要么内存爆，要么把 GPU 钉太久。我们能不能把 prompt 分块处理，或者用更聪明的方式去算 attention？ 主线 2 —— 让循环跑得快。 每次 forward 只产出一个 token，§9 已经把最大的成本指出来了：朴素的循环大部分时候都在重做它已经做过的事。这条主线下的文章，主题是别再重做、把 forward 摊给多个用户、调度谁什么时候跑。\n别再重做。 §6 已经把那条性质摆好了：旧 token 的表示一旦算出来就不再变。所以应该可以把它存下来下次复用，而不是每次都重算。这个状态得放在某个地方 —— 放哪？它有多大？随着对话变长它怎么涨？ 同时来很多用户。 真实的 serving 引擎会同时跑很多 prompt，长度都不一样，结束时间也都不一样。怎么让它们共享一次 forward 又不被 padding 拖累？当一些用户还在 token 1、另一些已经到 token 1000 的时候，scheduler 怎么让所有人都在前进？(Article 04 就是这条线的开头。) 处理一段长 prompt 跟\u0026quot;再多生成一个 token\u0026quot;完全不像。 §9 里每次调用的成本，在你是从头处理一段长输入还是只追加一个 output token 的时候，差得非常远。它们的瓶颈落在 GPU 的不同部位。引擎也许就该把它们当作两种不同的 workload —— 甚至放到不同的机器上去跑。 §1–§7 里那个模型，是这两条主线都在讨论的对象；§8 那个循环，是它们都想让它在规模化场景下跑得动的东西。系列后面的文章，就是把这些问题一道一道挑出来、逐个回答。\n","permalink":"https://wgzesg.github.io/llm_stories/zh/posts/01-llm-end-to-end/","summary":"三个 zoom level —— 整张网络、一个 transformer block 的内部、生成一段文字的完整循环。够你在脑子里搭起一个 LLM 的样貌，也够你开始问对的问题。","title":"LLM 从头到尾走一遍"},{"content":"这不是一份教程。这是一段在你脑子里搭 mental model 的旅程 —— 每读完一节，你会忍不住想 \u0026ldquo;哦，原来就这么回事？\u0026quot;。读到最后，tensor parallelism 不再是一个工程上的奇技淫巧。在那个场景下，它会变成 —— 你能想到的最自然不过的两个选择。\n不写矩阵公式。只讲 shape 和故事。\n1. 关于\u0026quot;输入\u0026rdquo;，你只需要这一张图 先别把 token 当成\u0026quot;词\u0026quot;。在模型眼里，一个 token 就是一行数字 —— d 个数。要装一点的话，叫 \u0026ldquo;feature vector\u0026rdquo;。\n一整句话就是这堆行的堆叠：\nToken 1 → [ f1 f2 f3 ... fd ] Token 2 → [ f1 f2 f3 ... fd ] Token 3 → [ f1 f2 f3 ... fd ] ... Token n → [ f1 f2 f3 ... fd ] 就这。n 个 token，每个都活在 d 维空间里。把这张图记牢 —— 后面所有东西都建立在它之上。\n2. 这个矩阵到底从哪冒出来的 在我们抽象地玩 weight matrix 之前，先在真实的 LLM 推理场景里找一个具体的落脚点 —— 这样 shape 才有实感，而不是一串空气。\nLLM 处理你的 prompt 时，第一个大阶段叫 prefill：把 n 个 prompt token 一次性全部塞进网络。（一个一个吐回答 token 是后面的 decode 阶段。）而 prefill 内部第一个计算，就是 attention 里的 QKV projection —— 每个 token（长度 d）要变成一个 query、一个 key、一个 value（每个长度 k）。\n把 token 堆成 section 1 那张 n × d 表，整个 QKV 这一步（这里只画 Q）就是一次矩阵乘法：\n[ n × d ] @ [ d × k ] = [ n × k ] tokens weight 每个 token matrix 的 query shape 就长这样。先停一下，把这个问题嚼一会儿：\n这个 matmul 到底在做什么？ 一张 n × d 的 token 表去乘一个 d × k 的 weight matrix，到底是个什么意思？\n\u0026ldquo;算出 query 啊\u0026quot;是无聊的回答。有意思的问题是 —— 这个 d × k 矩阵内部到底发生了什么。 这里有两个完全不同的故事可讲。每个故事都会悄悄给你一种把工作切到多 GPU 上的方式。\n3. 同一个 weight matrix，两种讲法 linear layer 把一个 token（长度 d）变成长度 k 的东西。干这活的是一个 shape 为 d × k 的 weight matrix。\n一个超级重要的题外话。 我说 \u0026ldquo;linear layer\u0026rdquo; 的时候，不是指网络里某一个特定的 block。我是指 transformer 里每一个 matmul：\nattention 里的 Q, K, V projection —— 每个都是把 token 变成 query / key / value 的 d × k 矩阵 attention output projection FFN up-projection (d → 4d) 和 FFN down-projection (4d → d) 甚至最末尾的 unembedding 它们都是同一种 shape 的运算：token 进，矩阵乘，token 出。所以下面要讲的\u0026quot;两种视角\u0026rdquo;，以及由此引出的两种并行策略，对所有这些都适用。一个看明白了，整个 transformer 的 matmul 都看明白了。\n接下来就是有意思的部分了：一个 d × k 矩阵可以用两种方式读 —— 一列一列地读或一行一行地读。同样的数字，同样的乘法，但脑子里是两幅完全不同的画面。两个我们都会走一遍。\nStory A —— 一列一列地读（一排 fx） 别再把 weight matrix 当成一堆数字组成的格子。镜头拉远。每一列都是一个独立的小函数 —— 给它一个 token，它返回一个数字。我们就把每一列叫做 fx（feature extractor 的缩写），然后把整个 weight matrix 画成一排 k 个 fx：\nweight = [ fx1 fx2 fx3 ... fxk ] 整个矩阵就这样。不再是数字 —— 是 fx。每一个都是它自己独立的小黑盒。\n每个 fx 怎么把 token 变成一个数字？其实就是和它那一列的 d 个权重做一个内积。但说实话 —— 为了建立 intuition，你不用关心。它就是 \u0026ldquo;fxi 看一眼 token，给出一个分数\u0026rdquo;。\n那这一层 layer 作用在 token 上，就只是：让 token 从这一排 fx 走过一遍，把每个 fx 吐出来的数字接住。\ntoken ⇒ [ fx1 fx2 ... fxk ] ↓ ↓ ↓ [ fx1(token), fx2(token), ..., fxk(token) ] 一个 token 走过 k 个小评委（feature extractor），每个吼出一个数字，你按顺序收集起来。output 长度就是 k。完事。\nStory B —— 一行一行地读（一摞 basis vector） 现在把同一个矩阵平铺。它有 d 行，每行长度 k：\nRow 1 → [ r1 r2 r3 ... rk ] Row 2 → [ r1 r2 r3 ... rk ] ... Row d → [ r1 r2 r3 ... rk ] 每一行都是 output space（长度 k）里的一个 basis vector。而 token 的 d 个 feature 就是 coefficient —— 告诉你每一行该掺多少进去。\noutput = f1 · Row1 + f2 · Row2 + ... + fd · Rowd 这一层 layer 用这种讲法就是：拿 token 的 d 个 feature 当配方，把 d 个 row vector 线性组合成一个 output vector。\n\u0026ldquo;等等，怎么……\u0026rdquo; 的瞬间 两种讲法描述的是完全相同的乘法。同样的数字进去，同样的数字出来。但你脑子里的画面截然不同：\nStory A (column) Story B (row) 多个独立的fx 一次大的线性组合 \u0026ldquo;从 token 里提取k 个 feature\u0026rdquo; \u0026ldquo;把d 个 row vector 混成 output\u0026rdquo; output 是被收集起来的 output 是被加起来的 这个二元性不是个 trivia —— 它就是 tensor parallelism 的种子。矩阵能怎么读，就能怎么切到 GPU 上。\n4. 现在你有两块 GPU。最自然的事是什么？ 一个矩阵，两块 GPU。盯着这个矩阵看。其实你能在它上面画的\u0026quot;自然的\u0026quot;切线只有两条：要么竖着切，要么横着切。\nSection 3 刚刚告诉你，每条线意味着什么。\n5. Strategy A —— 切 fx (Column Parallel) 把 Story A 当真。weight matrix 就是一排 k 个黑盒 fx。把它切到两块 GPU 上 —— 字面意义上 —— 就是在这排上画一条竖线：\nweight = [ fx1 ... fx(k/2) ‖ fx(k/2+1) ... fxk ] ↑ ↑ └──── GPU 1 ───┘ └──── GPU 2 ───┘ 每块 GPU 都看到完整的 token。 它只是跑它自己那一半 fx。\nGPU 1 → [ fx1(token), ..., fx(k/2)(token) ] GPU 2 → [ fx(k/2+1)(token), ..., fxk(token) ] 要拼出最终 output，直接拼起来就行：\noutput = [ GPU1 那半 | GPU2 那半 ] 完事。中间不需要求和，不需要同步。每块 GPU 在跑不同的 fx，输入是同一个，结果就这么紧挨着放。\n通信成本：便宜。拼接基本不要钱。\n6. Strategy B —— 切 row (Row Parallel) 把 Story B 当真。weight matrix 是一摞 d 个 basis vector row。把它切到两块 GPU 上 —— 字面意义上 —— 就是横着画一条线：\nweight = [ Row 1 ] ┐ [ Row 2 ] │ GPU 1 (配 feature 1..d/2) [ ... ] │ [ Row(d/2) ] ┘ ───────────────────────── [ Row(d/2+1) ] ┐ [ ... ] │ GPU 2 (配 feature d/2+1..d) [ Row d ] ┘ 但这里有个微妙的事：每一行都要乘上对应的 token feature（Row i 配 f_i）。所以切了 row，自动就切了输入 —— GPU 1 永远只需要 f_1..f_(d/2)，GPU 2 只需要剩下那一半。\n每块 GPU 只看到 token 的一半。 它产出的 output 是个长度为 k 的 vector —— 但只是总和的一部分。\nGPU 1 → partial output (它的 row 乘它的 feature) GPU 2 → partial output (它的 row 乘它的 feature) 要拼出最终 output，得加起来：\noutput = GPU1 的 partial + GPU2 的 partial 这次不能直接拼接 —— 两块 GPU 各自产出长度为 k 的 vector，需要逐元素相加。这个加法必须跨 GPU 完成。（这就是 TP 论文里那个 \u0026ldquo;all-reduce\u0026rdquo;。）\n通信成本：贵。这一层每次 forward 都要跨 GPU 做一次求和。\n7. 两种策略，并排对比 切 column (A) 切 row (B) 基于哪个故事 \u0026ldquo;一排fx\u0026rdquo; \u0026ldquo;row 的加权组合\u0026rdquo; 每块 GPU 拥有什么 一部分fx 一部分 row + 配套 feature 每块 GPU 看到的输入 完整的 只有一部分 输出怎么合 拼接 求和 (all-reduce) 通信开销 便宜 贵 同一个矩阵。两个故事。两种切法。这就是全部游戏规则。\n8. Multi-Head Attention：切口本来就在那 把这套东西用到 transformer 里一个真实的部件上 —— attention 里的 QKV projection —— 看着 column-parallel TP 怎么免费掉出来。\n场景 Q（K, V 同理）projection 把每个 token（长度 d）变成一个长度 k 的 query vector。但 k 不是一个随便的数字 —— 它是有结构的：\nk = h × d_head\n其中 h 是 head 数量，d_head 是 每个 head 的维度。\n所以我们那一排 fx 是有组织的。每 d_head 个相邻的 fx 分一组，每一组就叫一个 head：\nW_Q = [ fx1 ... fx(dh) │ fx(dh+1) ... fx(2·dh) │ ... │ fx((h-1)·dh+1) ... fxk ] └── Head 1 ────┘ └──── Head 2 ─────────┘ └──── Head h ──────────┘ Head 1 的 fx 产出 Head 1 的 query vector，Head 2 的 fx 产出 Head 2 的，以此类推。同一个矩阵，同一排 fx —— 只是分了组。\n实现层面的小注。 实际代码里这是一个大 matmul，shape [d, h × d_head]，不是 h 个小 matmul —— 一个大的矩阵乘法在 GPU 上比一堆小的快得多得多。\u0026ldquo;h 个 head\u0026rdquo; 这个结构活在每一列代表什么这件事里，不在矩阵的数量里。（很多生产环境的代码会更激进 —— 把 Q, K, V 三个 projection fuse 成一个 [d, 3 × h × d_head] 的大 matmul，一次算完。）数学上 —— 包括训练时 —— 完全没区别。head 这个结构纯粹是个逻辑上的分组。\n为什么 head 让 column-parallel 显得理所当然 下面是高潮部分。\nattention 内部的计算里，head 之间是互不相干的。Head 1 的 attention 只让 Head 1 的 query 和 Head 1 的 key 玩，Head 2 自己玩自己的，永远不互相看。所有 head 真正混在一起，要等到最后一个独立的矩阵 —— output projection。\n所以如果你反正都要把 fx 按 column 切到 GPU 上（也就是 section 5 里的 Strategy A）……就沿着 head 的边界切：\nW_Q = [ Head 1 │ Head 2 │ Head 3 │ Head 4 ] ↑ ↑ ↑ ↑ └── GPU 1 ─┘ └── GPU 2 ──┘ 每块 GPU 拥有几个 head。它独立完成自己那几个 head 的 Q, K, V projection 和 attention。整个 attention 期间零通信。 每块 GPU 都在跑自己的私有 mini-attention。\n真正的 aha multi-head attention 不是为 tensor parallelism 设计的。它的初衷是让不同的 head 学到关注输入里不同的关系模式 —— 这是个建模层面的选择，不是系统层面的。\n但 TP 出现的时候，它走过来一看：哎，attention 早就被预先切成一块块独立的 \u0026ldquo;head\u0026rdquo; 了。 它根本不需要发明任何东西。它只是顺着已经存在的切口把工作分了。\n这就是真实 transformer 里最干净的 column-parallel TP 案例 —— 它直接从 Story A 掉出来。matrix 是一排 fx，fx 按 head 分组，head 之间互相独立 —— 所以沿着 head 边界切。一个 mental model 用到底。\n","permalink":"https://wgzesg.github.io/llm_stories/zh/posts/02-tensor-parallelism-mental-model/","summary":"把 weight matrix 用两种方式读，就有两种把它切到多 GPU 上的方法。从 transformer prefill 里的一次 matmul，推出 tensor parallelism 的整套心智模型。","title":"Tensor Parallelism 心智模型：从零搭起"},{"content":"Article 02 留给我们两种把一次 matmul 切上两张 GPU 的方式。直接看它们在做什么，比记 paper 里那些名字省事得多。先把两种摆一起对一下：\nStrategy A —— 切 fx Strategy B —— 切行 切的是矩阵的什么 列（每列是一个 fx） 行（每行是一个基向量） 每张 GPU 拿到的输入 每个 token 的完整输入向量 每个 token 输入特征的一半 每张 GPU 算出的输出 每个 token 输出特征的一半 每个 token 完整输出的一个 partial sum 怎么合起来 concatenate（免费） all-reduce（一次通信） 别名 column-parallel row-parallel 每一列最精简的读法：A = \u0026ldquo;full in, half out\u0026rdquo;，B = \u0026ldquo;half in, sum out\u0026rdquo;。后面所有东西都建立在这两行字上。\n但一个真实的 transformer block 不是单次 matmul，是 四次，再加一些 pointwise 的胶水。下一个自然冒出来的问题就是：一整个 block 该怎么切上两张 GPU？\n第一反应有一个非常自然的方案，差点就跑通了。我们先把它搭出来，看清楚它在哪儿坏掉，再让\u0026quot;补这个洞\u0026quot;的过程把我们带到 Megatron 那个经典 pattern。整篇都用一组小数字盯着 shape 走，免得讨论变空气。\n1. 起手：一组装得下的小数 两张 GPU，叫 G1 和 G2。一个小 batch 配一个小模型：\n值 batch n 4 个 token model dim d 512 heads h 8 per-head dim d_head 64 attention dim k = h · d_head 512 FFN hidden 4d = 2048 每个 token 是一行 512 个数。整个 batch 就是 [n × d] = [4 × 512]。\n把一个 transformer block 拍平：\n← input: [4 × 512] │ LayerNorm │ QKV projection d → 3k ← matmul weight [d × 3k] = [512 × 1536] │ attention (Q 和 K 互相混；不是一次新的 matmul) │ output projection k → d ← matmul weight [k × d] = [512 × 512] │ + residual │ LayerNorm │ FFN up-projection d → 4d ← matmul weight [d × 4d] = [512 × 2048] │ activation (GeLU) (pointwise) │ FFN down-projection 4d → d ← matmul weight [4d × d] = [2048 × 512] │ + residual │ 四次 matmul，加一些胶水。\n顺嘴说一下 pointwise 这些胶水：为什么两张 GPU 都要算一遍？ LayerNorm、activation、residual add 这些都是 pointwise 的。它们不在乎数据怎么分布在 GPU 上，只要本地有它要的那份就行。在 TP 里我们偷个懒：只要某份数据完整地在两张 GPU 上都有，就让两张 GPU 各跑一遍这个 pointwise op。同样输入、同样输出，重复算。为什么不让一张 GPU 算完再 broadcast？因为 瓶颈一直是通信，不是算力。pointwise 在 GPU 上对几千个数做一遍，几乎不花时间；跨 GPU 发数据是真金白银的 latency 和 bandwidth。多花点便宜的算力比省那点通信划算。把这条记着 —— 后面的 trace 表里你会看到每个 LN 和 residual 步骤都标着 \u0026ldquo;(冗余)\u0026quot;，就是这个意思。\n所以这个 block 整个 TP 故事，全在那四次 matmul 上。两张 GPU，四个 cut 要做。开搞。\n2. v1 —— 把 Strategy A（full → half）用到每个 matmul 上 最自然的第一步是什么？回到 article 02，Strategy A 长这样：\n便宜的那种 cut（concatenate，matmul 内部不需要 all-reduce）； 用在 QKV 上正好落在 head 边界：k = 8 · 64 = 512，每张 GPU 256，正好一人 4 个 head； 而且\u0026quot;完整输入进，半个输出出\u0026quot;也是更好脑补的那一版。 那就把 A 一路用到底，四次 matmul 全是 A。一步一步走过这个 block，盯着每张 GPU 手上有什么 —— 它的 weight 切片、输入、输出：\nStepGPU 1GPU 2 input [4×512] 完整 [4×512] 完整 LayerNorm (冗余) in [4×512] → out [4×512] in [4×512] → out [4×512] QKV proj (A) W [512×768] (heads 1–4)\nin [4×512] → out [4×768]\n= heads 1–4 的 Q+K+V，每份 [4×256] W [512×768] (heads 5–8)\nin [4×512] → out [4×768]\n= heads 5–8 的 Q+K+V，每份 [4×256] attention heads 1–4\nin [4×768] → out [4×256] heads 5–8\nin [4×768] → out [4×256] ★ GATHER #1 —— output proj 要的是完整 k=512，每张 GPU 只有 256，凑出 [4×512] output proj (A) W [512×256]\nin [4×512] → out [4×256] W [512×256]\nin [4×512] → out [4×256] ★ GATHER #2 —— residual 要完整 d=512，输出只有半个 d=256，凑出 [4×512] + residual [4×512] → [4×512] [4×512] → [4×512] LayerNorm (冗余) in [4×512] → out [4×512] in [4×512] → out [4×512] FFN-up (A) W [512×1024]\nin [4×512] → out [4×1024] W [512×1024]\nin [4×512] → out [4×1024] activation (pointwise) [4×1024] → [4×1024] [4×1024] → [4×1024] ★ GATHER #3 —— FFN-down 要完整 4d=2048，每张 GPU 只有 1024，凑出 [4×2048] FFN-down (A) W [2048×256]\nin [4×2048] → out [4×256] W [2048×256]\nin [4×2048] → out [4×256] ★ GATHER #4 —— residual 要完整 d=512，输出只有半个 d=256，凑出 [4×512] + residual [4×512] → [4×512] [4×512] → [4×512] 每个 block 四次跨 GPU 的 gather。\n其中两次发生在下一个 A 风格的 matmul 之前 —— 它要的是完整输入。另外两次在 residual add 之前 —— 它要完整向量，但我们刚算出来的是半个。本质都是同一个原因：A 产出半个输出，下游几乎所有东西要的都是完整输入。\n3. v1 的代价 跨 GPU 通信是分布式计算里 慢 的那一环。整个 TP 设计的目标，就是把它压到尽可能少。v1 的现状是：几乎每一个需要完整 feature 的算子前面，都得付一次 gather。\n一个 32 层的模型，光一次 forward 就 ~130 次跨 GPU 通信。太多了。\n问题就变成了：\n能不能把 gather 省掉？\n每次 gather 的存在都是一个原因：下一个算子要完整向量，但 Strategy A 给的是半个。我们真正想要的，是一个 愿意 直接吃半个输出的 matmul。\narticle 02 已经把它递到我们手上了。\n4. v2 —— 让 Strategy A 配上 Strategy B（half → sum） 换一个角度看这两种 strategy：\nStrategy A 输出 一个 half。 Strategy B 输入 一个 half。 形状一样。 A 的输出正好是 B 想要的输入。两者咬合上，中间一点通信都不需要。\n把 v1 里的 \u0026ldquo;A → gather → A\u0026rdquo; 换成 \u0026ldquo;A → B\u0026rdquo;。B 直接吃下半个输出，而通信代价只剩下 B 末尾那一次 —— 把 partial sum 加成完整输出的 all-reduce，给后面的 residual 和 LN 用。\n把这个套路用到整个 block，每个 A 配一个 B：\nStepGPU 1GPU 2 input [4×512] 完整 [4×512] 完整 LayerNorm (冗余) in [4×512] → out [4×512] in [4×512] → out [4×512] QKV proj (A) W [512×768] (heads 1–4)\nin [4×512] → out [4×768]\n= heads 1–4 的 Q+K+V，每份 [4×256] W [512×768] (heads 5–8)\nin [4×512] → out [4×768]\n= heads 5–8 的 Q+K+V，每份 [4×256] attention heads 1–4\nin [4×768] → out [4×256] heads 5–8\nin [4×768] → out [4×256] output proj (B) W [256×512]\nin [4×256] → out [4×512] (部分和) W [256×512]\nin [4×256] → out [4×512] (部分和) ★ ALL-REDUCE #1 —— 把两张 GPU 上的两个 [4×512] 部分和加起来，凑出完整 [4×512]（residual + LN 要用） + residual [4×512] → [4×512] [4×512] → [4×512] LayerNorm (冗余) in [4×512] → out [4×512] in [4×512] → out [4×512] FFN-up (A) W [512×1024]\nin [4×512] → out [4×1024] W [512×1024]\nin [4×512] → out [4×1024] activation (pointwise) [4×1024] → [4×1024] [4×1024] → [4×1024] FFN-down (B) W [1024×512]\nin [4×1024] → out [4×512] (部分和) W [1024×512]\nin [4×1024] → out [4×512] (部分和) ★ ALL-REDUCE #2 —— 把两张 GPU 上的两个 [4×512] 部分和加起来，凑出完整 [4×512]（residual + LN 要用） + residual [4×512] → [4×512] [4×512] → [4×512] 每个 block 两次 all-reduce。\n这就是 Megatron pattern。没人告诉我们答案，我们自己一步一步走进去的。\n5. 没料到的对偶性 article 02 把 A 和 B 当两种独立 strategy 介绍 —— 同一个矩阵的两种读法。但把它们并排放着看，盯着每种的输入输出形状：\nA 拿 完整输入，吐 半个输出。 B 拿 半个输入，吐 完整 sum 作为输出。 它们不是两种 strategy，是 同一次往返的两半。A 的输出 shape 就是 B 的输入 shape；B 的输出 shape（all-reduce 之后）就是 A 的输入 shape。你发明 A 的同时，其实就把 B 当成它的回程一起发明出来了。\n再回头看这个 block 在做什么：\nAttention 是 widen（QKV：d → k）后跟 narrow（output proj：k → d）。 FFN 也是 widen（d → 4d）后跟 narrow（4d → d）。 widen 那里输出 feature 多，正好分到多张 GPU —— A 派得上用场。narrow 那里输入 feature 多，正好把输入切到多张 GPU，输出小再加回来 —— B 派得上用场。\n这个 block 不是\u0026quot;碰巧\u0026quot;对 A→B 友好，是 结构上就对 A→B 友好：两对 widen-narrow，中间用 pointwise 粘起来。\u0026ldquo;Megatron pattern\u0026rdquo; 不是哪个人坐下来设计出来的算法，它就是唯一一个尊重架构本身做的事的通信 pattern。A、B 对偶和 widen-narrow 节奏，是同一件事讲两遍。\n提一句通信代价：一次 gather 和一次 all-reduce 在每张 GPU 上搬的数据量差不多（all-reduce 内部大致是 reduce-scatter 加 all-gather）。v1 每个 block 4 次 gather；v2 每个 block 2 次 all-reduce —— 通信砍一半，模型一行没改。\n6. 为什么 cut 必须落在 head 边界 v2 trace 偷偷假设了一件事：QKV 的 column cut 是沿着 head 边界 把 k = 512 切成两块各 256，每张 GPU 正好拿 4 整个 head。这个假设干的活比看起来多得多。试一下反事实就知道。\n想象一下 single-head attention —— 同样 k = 512，但只有一个 head，没有 head 结构。Strategy A 照旧用到 QKV 上：每张 GPU 拿到形状都是 [4 × 256] 的 Q, K, V。开始算 attention。\n第一步是 Q Kᵀ。每张 GPU 算 Q_half @ K_halfᵀ，得到一个 [4 × 4] 的矩阵 —— 但这个矩阵是它手上那 256 个 feature 的 partial sum。真正的 scores 得把两张 GPU 上的 partial 加起来才行。\n接下来麻烦了：下一步是 softmax。Softmax 是非线性的，本地算完再合不行 —— softmax(a) + softmax(b) ≠ softmax(a + b)。reduction 必须发生在 softmax 之前。也就是说，attention 中间会硬塞进一次同步：\n★ ALL-REDUCE，对 [n × n] 的 scores，发生在 softmax 之前。\n这是 v2 那两次之外的第三次 all-reduce。Megatron pattern 直接塌成 三 次同步，而新增的这次的 tensor 大小随 sequence length 平方增长 —— 你最不想让它变大的那个。\n修法不在算法层，在结构层：不要让 cut 跨过 head。每个 head 的 Q Kᵀ 必须完整地落在一张 GPU 上，partial sum 的问题就不会出现。multi-head attention 把这个白送给我们了：head 在定义上就互相独立，head 边界天然是切点；只要 h 能被 GPU 数整除，对 k = h · d_head 的 column split 就正好落在两个 head 之间。\n所以 multi-head 不是系统人捡了 modeler 的便宜。它是 v2 能存在的 结构性前提。cut 落到 head 内部，softmax 立刻逼出一次毁掉一切的同步；cut 落在 head 之间，非线性就保持本地。Megatron pattern 不是 碰巧 在 multi-head 架构上能跑 —— 它要求 multi-head 架构。\n7. 这一篇打开了哪些门 到这里，一个 block 跑在两张 GPU 上，每次 forward 两次 all-reduce。下一轮\u0026quot;等下，那……\u0026ldquo;的问题就开始了：\n如果 block 很多、GPU 也很多呢？ TP 切的是 block 内部。切 跨 block —— 把整个 block 摆到不同 GPU 上，让 microbatch 流过去 —— 是另一回事。Pipeline parallelism，下一篇见。 如果 FFN 被换成 expert 呢？ column-then-row 这一套对每个 expert 的 matmul 还是适用，但把 token 路由到正确的 expert 又引入一种新通信。MoE，再下一篇。 如果 batch 里 sequence 长度差异巨大呢？ 通信 pattern 不变，但 attention 那边要处理变长 sequence —— continuous batching 由此登场。 同样的语法。每条都自成一篇 walk-through。\n","permalink":"https://wgzesg.github.io/llm_stories/zh/posts/03-tp-through-a-full-block/","summary":"把 article 02 的两种 cut 方式拿到一整个 transformer block 上跑一遍，盯着每一步每张 GPU 上的 shape。先把一种 cut 用到所有 matmul 上 —— 通信爆炸，每个 block 四次 gather。再把两种 cut 配成一对，刚好对上 widen-narrow 的架构节奏，落到每个 block 两次 all-reduce。","title":"在一个 transformer block 中完整走完一遍 Tensor Parallelism"},{"content":"Article 03 把一个 transformer block 跑在两张 GPU 上，每层两次 all-reduce 就拿下了。但真实的 serving 系统不会只有一个用户 —— 一堆 prompt 同时打进来，长度还各不相同。一个 50 token 的 \u0026ldquo;现在几点了\u0026rdquo; 就坐在一份 5,000 token 的论文草稿旁边。\n这一篇要追两个问题：\n怎么把不等长的 request 高效地塞进一次 forward？ 最自然的做法 —— 把所有 prompt 都 pad 到最长那个，按固定 batch 跑 —— 在短的那些上浪费一大堆算力。肯定有更聪明的办法。 TP 这边要不要知道 batch 这件事？ 还是说 batching 的招数和切模型的故事可以彼此独立？ 我们继续沿用 article 03 的 setup，盯着每一层在面对多个 request 时做了什么。\n1. Setup 数字跟 article 03 一样：\n值 GPUs 2 张 (TP=2) layers 8 d (model dim) 512 h (heads) 8，每张 GPU 4 个 d_head 64 k = h · d_head 512 FFN hidden 4d = 2048 讨论里我们用两个具体的 request：request A 长度 10，request B 长度 30。\n这一篇明确握着三条假设：\n只考虑 prefill。 我们只算每个 request 的 prompt 那次 forward。逐 token 的 decode 还没来 —— 那是 article 05 的事。 每个 request 都装得下一个 batch。 一个 batch 装的是 ≥1 个整 request，从来不会装\u0026quot;半个\u0026quot;。article 06 用 chunked prefill 来放松这条。 暂时不引入 KV cache。 KV cache 是 decode 阶段让后面的 token 能看回之前 token 的那个东西。在只有 prefill 的世界里，我们算完输出直接发出去，没什么需要存下来的。KV cache 跟着 article 05 一起到。 这样把空间维度的故事讲干净。时间维度的故事（跨 iteration 的 continuous batching）是另一篇。\n2. 先看一个 request：N 只是个 tensor 维度 在两个 request 之前，先回忆一下 article 03 v2 里一个 request 长什么样。一个长度 N 的 prefill 走过这个 block，shape 就是 [N × 512]。从那篇的 trace 里可以看到：\n8 层 × 每层 2 次 all-reduce = 每次 forward 16 次 all-reduce。 每次 all-reduce 跨 GPU 搬的都是 [N × 512]。 值得停一下的是：N 只出现在 tensor shape 里，从来没出现在通信次数里。 不管 N=10 还是 N=10,000，你都做正好 16 次 all-reduce。区别只是每次搬的字节数多还是少。\n也就是说，给一个 request 加更多的 token 在通信成本上是\u0026quot;白送\u0026quot;的 —— 总字节数线性地涨，但同步事件的次数没增加。\n这是个不错的性质。下一个要问的是：当多出来的这些 token 来自不同 request 的时候，这条性质还能不能保住。\n3. 自然的做法和更聪明的做法 自然做法：pad 到最长。 A 和 B 摞起来变成 [2 × 30 × 512]。A 拿到 20 个 padding token，模型还是要照算。Linear 那边的浪费温和（matmul 大了 2×）。Attention 那边就重了 —— 每个 request 的 attention 是 O(L²) 的活，A 的 attention 要算 30² = 900 次（per head per layer），实际只需要 10² = 100 次。单 A 一个就多算了 9 倍，padding token 还对最终结果没贡献。\n更聪明的做法：flatten。 把 A 和 B 的 token 拼成一个 shape 为 [(10+30) × 512] = [40 × 512] 的 tensor。没有 padding，没有 batch 维度 —— 就是一长串 token。\n但问题立刻冒出来：这个 flatten 之后的 tensor 走过一整个 forward 的每一步，还能算出对的结果吗？ 有些步骤显然没问题。有些得想一想。一步一步走过去看看。\n4. 从头到尾把这个 block 走一遍 从输入 [40 × 512] 开始，把这个 block 的每一步过一遍。每一步都问：当输入里同时有多个 request 的 token 时，它能不能算出正确答案？\n步骤 它做什么 在 [40 × 512] 上？ LayerNorm 每行各自归一化 ✓ 直接没问题 QKV proj (linear) 跟共享的 W 做 matmul 需要分析 Attention 逐 request 的 sequence-mixing 需要分析 Output proj (linear) 跟共享的 W 做 matmul 需要分析 Residual add 每行做加法 ✓ 直接没问题 LayerNorm 每行各自归一化 ✓ 直接没问题 FFN-up (linear) 跟共享的 W 做 matmul 需要分析 Activation (GeLU) 每个元素的非线性 ✓ 直接没问题 FFN-down (linear) 跟共享的 W 做 matmul 需要分析 Residual add 每行做加法 ✓ 直接没问题 一半的步骤一上来就打钩了。Pointwise 算子 —— LayerNorm、GeLU、residual add —— 处理每行都是独立的。第 i 行属于 request A 还是 request B，对它们来说是不可见的。它们逐 token，互不混合。所以 batching 是白送的。\n剩下五步要细看：四个 linear matmul（QKV、output、FFN-up、FFN-down）外加一个 attention block。\n但有个便利：四个 linear matmul 的结构完全一样 —— Y = X @ W，W 是被所有行共享的。所以只要弄清楚 一个 linear 在 flatten batching 下的行为，四个全都跟着确定下来。Attention 在每层只有一个。\n整个 batching 的问题就缩成两条：\n一个 linear layer 在 [40 × 512] 上算出来对吗？ Attention 在 [40 × 512] 上算出来对吗？ §5 处理 linear，§6 处理 attention。这两条解决，整个 block 就解决了。\n5. Linear layer：简单的那一半 回到 article 02 是怎么看 linear layer 的。weight matrix 是一排小小的 feature extractor —— 每个 fx 是它自己一个不透明的小函数，吃一个 token 的 d 维 feature vector，吐一个数。一个 k 维输出的 linear layer，就是 k 个这样的 extractor 并排站着，对同一个 token 一齐工作。\ntoken ⇒ [ fx1 fx2 fx3 ... fxk ] ⇒ [ fx1(token), fx2(token), ..., fxk(token) ] 值得停一下的事：每个 fx 看的是 一个 token，吐出 一个数。它不偷瞄下一个 token，也不偷瞄上一个 token。它不知道这个 token 来自哪一个对话。整个 math 里就没有让 request 边界进来的入口，因为它一次只看一个 token。\n那把一个 flat tensor [40 × 512] —— 40 个 token 摞起来 —— 喂给它，它就把每个 fx 在每个 token 上跑一遍。40 个 token、每个 k 个 extractor，填满一个 [40 × k] 的输出。前 10 行是 request A、后 30 行是 request B 这件事，对这次操作来说根本看不到；它从一开始就没有混淆它们的可能性。\n这就是 linear layer 在 batching 下白送的根本原因。它们不是\u0026quot;奇迹般地\u0026quot;能 batch —— 它们本来就是逐 token 的。我们只是让它跑了更多个 token 而已。\n**TP=2 的情况：**跟 article 03 一样没变。fx 还是按 head 切到两张 GPU 上，每张 GPU 拥有一半：\nG1 在 [40 × 512] 上跑 heads 1–4 的 fx → [40 × 768] G2 在 [40 × 512] 上跑 heads 5–8 的 fx → [40 × 768] all-reduce 的 shape 从 [N × 512] 变成 [40 × 512]，但 all-reduce 的次数没变。同样的通信 pattern，每次搬的字节多了。\n而且这个论证对 block 里所有四个 linear matmul 都适用 —— QKV、output、FFN-up、FFN-down —— 所有 linear 全都搞定了。 只剩一步。\n6. Attention：难的那一半 为什么 attention 不一样？因为 attention 是 sequence-mixing。每个 token 的输出依赖 sequence 里 所有 token，不只是它自己那一行：\nout[i, :] = softmax( Q[i, :] @ K.T / √d_head ) @ V 这里的 K.T 和 V 是横扫整个 sequence 的。如果 K 和 V 来自一个同时装着 A 和 B token 的 tensor，那 A 里第 i 个 token 默认就会去 attend B 的 token —— 反过来也一样。math 能跑通，但答案是错的：A 的输出会跟 B 的 key、value 混在一起，这不是模型训练时学到的东西。\n我们需要一种办法：让 request A 的 attention 严格只在 A 的 token 范围内做，B 的也严格只在 B 的范围内 —— 但底下的 flat tensor 还是共享的。\n6.1 自然做法 —— 先算，再 mask 最直接的修法：把整个 [40 × 40] 的 attention matrix 当成 40 个 token 是一个 sequence 一样算出来，然后把跨 request 的那些位置 mask 掉（softmax 之前设成 -∞，让它们贡献为零）。\n这个 flat token buffer 长这样：\nflat token tensor: [40 tokens × 512] each row is one token of d=512 features request A 10 tokens request B 30 tokens row 0 row 10 row 40 cu_seqlens = [0, 10, 40] 完整的 attention matrix，跨 request 的 block 被 mask 掉之后：\nnaive: compute full 40×40, mask cross-request blocks keys 0..9 keys 10..39 queries 0..9 queries 10..39 A → A 10 × 10 masked to −∞ 10 × 30 masked 30 × 10 B → B 30 × 30 computed: 1600 useful: 1000 wasted: 600 跑是跑得通，但太浪费了。off-diagonal 那两块 —— 10 × 30 和 30 × 10，加起来 600 个位置 —— 算出来立刻就被丢掉。一旦并发的 request 多起来情况只会更糟：R 个长度都是 L 的 request，你算了 (RL)²，但只用得上 R · L²。跨 request 的活按 R² 涨，有用的活只按 R 涨。serving 系统里 R 轻易就能上百，这种做法根本撑不住。\n6.2 Varlen 的想法 —— 直接跳过，不要 mask 不要先算再 mask，干脆 只 算 diagonal 的那些 block。一个 request 一个 request 地循环，每次在它在 flat buffer 里那一段上跑普通 attention：\nvarlen: compute only the diagonal blocks keys 0..9 keys 10..39 queries 0..9 queries 10..39 A → A 10 × 10 B → B 30 × 30 (not computed) (not computed) computed: 1000 — no waste 这就是 variable-length attention kernel —— 简称 varlen。它吃一个 flat tensor 加一组 request 边界（cu_seqlens，cumulative sequence lengths），然后一个 request 一个 request 地走：\n# cu_seqlens = [0, 10, 40] # request A 占 [0,10)，B 占 [10,40) for i in range(num_requests): s, e = cu_seqlens[i], cu_seqlens[i+1] Q_i = Q[s:e] K_i = K[s:e] V_i = V[s:e] scores_i = (Q_i @ K_i.T) / sqrt(d_head) # L_i × L_i probs_i = softmax(scores_i + causal_mask_i) out[s:e] = probs_i @ V_i # 写回 flat buffer 把这个循环画出来，就是顺着 flat 的 Q、K、V stack 从上往下走：\nvarlen walks the flat Q, K, V stacks request-by-request Q K V A [0:10] A [0:10] A [0:10] B [10:40] B [10:40] B [10:40] i = 0 read slice [0:10] i = 1 read slice [10:40] step 1 — compute scores = Q_A @ K_A.T probs = softmax(scores) out[0:10] = probs @ V_A step 2 — compute scores = Q_B @ K_B.T probs = softmax(scores) out[10:40] = probs @ V_B flat Q, K, V tensors in HBM — varlen kernel slices [s:e] for each request, top to bottom 三件值得注意的事：\n跨 request 的 block 不是被 mask 掉，是根本没算。kernel 整块跳过去。 每次循环里的 score matrix 大小恰好是那个 request 的 —— [L_i × L_i]，不是 [40 × 40]。所以 score 那块的内存占用一直很小。 flat 输出 buffer 是这样填的：每个 request 算出来的 attention output 被写回到它自己那一段。 cu_seqlens 是 kernel 唯一需要知道的\u0026quot;关于 request 的\u0026quot;信息。剩下的就是 flat tensor 的切片操作。\n（实际跑的 kernel 不会真的在 Python 里循环 —— 它把循环放进 GPU、一次 launch 全部搞定，不会每个 request 多付一次 launch overhead。数学内容跟上面这段循环一样，优化版只是表达得更高效。高性能 attention kernel 是后面专门一篇的事。）\n6.3 在 TP=2 下 每张 GPU 还是各自管 article 03 那 4 个 head。varlen kernel 跑在每张 GPU 自己本地的 Q、K、V 上 —— 对它的那 4 个 head，对所有 request 的 token。G1 不需要知道 G2 在 attention 里干什么；G2 的 head 是 G2 的事。article 03 靠的那条\u0026quot;head 之间互相独立\u0026quot;的性质，在这个循环里照样成立。没有新增任何通信。\nlinear（§5）和 attention（§6）都搞定了，整个 block 在 batch 这件事上就都对了。\n7. 退一步看：TP 完全没动到 值得停一下的\u0026quot;意外之喜\u0026quot;。看一下整次 batched forward 里 TP 看到的东西：\n一个 shape 是 [tokens × hidden] 的 tensor 流过每一层。 weight 按 head 切。 all-reduce 在 [tokens × hidden] 的 partial sum 上做。 每个 block 16 次同步事件，跟 article 03 一模一样。 TP 完全没看到 request 边界。 那个 flat tensor 在 TP 眼里长的样子，跟 40 个 token 来自 1 个 request 还是 50 个 request，没有任何区别。request 边界只在一个地方进入 —— varlen attention kernel 里的 cu_seqlens 参数 —— 而这个参数完全在每张 GPU 自己的本地切片里用，没有引发任何通信。\n也就是说，request batching 和 TP 是两条互相不挡道的轴，唯一相交的地方在 attention kernel 内部：\nTP 回答的是：模型怎么切到多张 GPU 上？ Request batching 回答的是：token 怎么打包进一次 forward？ 这两个问题彼此不约束。我们没有为这一点做过设计 —— 它是从下面两条本来就成立的事实里掉出来的：\nlinear layer 是逐 token 的（所以即便在一张 GPU 上也看不到 request 边界）。 multi-head attention 的 head 互相独立（所以每张 GPU 自己的 per-head varlen 循环根本不需要跟其他 GPU 说话）。 article 03 的收尾说，multi-head attention 是 modeler 留给做系统的人的一份礼物，让 TP 通信免费。这里能看到这份礼物又往前走了一层：让 TP 通信免费的那条 head 独立性，同样让 request batching 通信免费。两件本来无关的招数因为同一条架构性质，刚好可以白嫖式地组合起来。\n8. 计算量分析 把 request flatten 起来之后，时间到底花在哪里，老实讲几句。\nLinear layer 看起来很爽。一份 weight 从 HBM 读上来，被摊到 flat tensor 的 (N+M) 个 token 上摊薄。塞进去的 token 越多，GPU 越接近 compute 的峰值。这就是为什么激进的 prefill batching 在 throughput 上是稳赚不赔的。\nAttention 这边复杂一些。每个 request 的 Q_i K_i.T 是它自己那次 matmul，没法像 linear 那样把所有 request 融成一次大 GEMM。现代 varlen kernel 把循环放进 GPU 一次 launch，所以不会按 request 数付 launch overhead。但每个 request 的 attention 还是 O(L_i²)，瓶颈长什么样很大程度上取决于 request 长度的分布。\n想象两种 batch，总 token 数一样：\n1 个 request × 1,000 token —— attention 的活是 1 × 1000² = 10⁶（per head per layer）。整个方块就是一整块。Attention 主导整次 forward。 10 个 request × 每个 100 token —— attention 的活是 10 × 100² = 10⁵，少了 10 倍。Linear 主导。 其实就是 §6.2 那张 varlen 方块图，被推到了两个极端。把全部 1,000 个 token 摆在 attention 矩阵的一条边上，varlen 真正会算的只有 per-request 那些对角块，剩下的都是跨 request 的格子，直接跳过：\n总 token 数都是 1,000，attention 的活差很多 彩色 = per-request 真正算的部分；斜纹 = 跨 request，varlen 直接跳过 1 个 request × 1,000 token 1 × 1000² = 10⁶ 没有跨 request 的浪费，attention 主导 10 个 request × 每个 100 token 10 × 100² = 10⁵ 大部分方块都跳过；linear 主导 外框一样大，总 token 数一样，但彩色那一块——真正会算的部分——从左到右少了 10 倍。attention 的 L² scaling 意味着：长上下文 batch 是 attention-compute-bound；短 request 多的 batch 是 linear-bandwidth-bound。flatten 这套招在两种 regime 下都一样，但瓶颈换了位置。\n这也是 decode 那一篇的预告：当每个\u0026quot;request\u0026quot;一次只生成一个 token 时，per-request 的 Q_i K_i.T 退化成一个 1 × L_kv 的向量乘上一个 L_kv × d_head 的矩阵。arithmetic intensity 掉到 ~1，per-request 的 matmul 不再\u0026quot;够大\u0026quot;，整次 forward 变成被 weight 的 bandwidth 卡住的状态。完全不一样的优化目标 —— 这就是为什么 decode 自成一篇。\n9. 这一篇打开了哪些门 到这里，多个并发 prefill request 在 TP 模型上要怎么跑，我们已经有一套方案了：把 token flatten 成一个 tensor，每个 linear layer 都是一次大 matmul，每个 attention block 都是一次 varlen attention。模型本身的 TP 通信 pattern 不变。naive padding 的浪费消失了。每个 request 拿到的算力都是它真正需要的那么多 —— 不多不少。\n下面三个跟进的问题各自值得一篇：\n如果一个 request 要生成很多个输出 token 呢？ Prefill 是一个 prompt 一次过完。Decode 多了一个逐 token 的阶段，瓶颈完全不同，还多了一个新结构（KV cache）来记住之前的 token。Article 05 —— decode 和跨 iteration 的 continuous batching。 如果某个 request 长到一个 batch 装不下呢？ \u0026ldquo;每个 request 都装得下\u0026quot;这条假设有时候就是会破。修法是 chunked prefill —— 把 prompt 切成一段一段处理，沿路把 KV cache 一段段建起来。Article 06。 varlen attention kernel 在 GPU 上到底怎么跑得快？ 我们整篇用的是 naive 的 attention 数学。高性能那一版（FlashAttention）干脆不让 score matrix 实例化出来，用 tiled online-softmax 的递推。这是 kernel 层级的深入，值得专门一篇，会在系列后面出现。 每次都是同一个语法：从这一篇拿走一条假设，放掉，看看会发生什么。\n","permalink":"https://wgzesg.github.io/llm_stories/zh/posts/04-batching-many-requests/","summary":"很多个用户同时打过来，prompt 长度还都不一样。把一整个 transformer block 拿到一个 flatten 起来的多 request tensor 上跑一遍，看哪些 layer 是白送、哪些得真动手 —— 顺便看一下 TP 这边到底要不要改。","title":"一次 forward 怎么塞下很多个 request"},{"content":"Article 04 把\u0026quot;很多个 prefill 塞进一次 forward\u0026quot;这件事干净地解决掉了。但 prefill 只是一个 request 一生中的前半段。prompt 嚼完之后，request 进入 decode 阶段 —— 一次产出一个 token，有时跑上几百步，直到撞到 EOS 才停。真实的 serving 引擎看不到那种干净整齐的 prefill 批次；它看到的是一锅乱炖：刚到的 prompt、跑到一半的 decode、马上要结束的 request，全在同一时刻共用同一张 GPU。\n这一篇要走进这片混乱。生成流程的基础和 KV cache 的机制我们直接借用 Article 01 里讲过的，假设你已经熟悉。\n（先把一个名字钉在脑子里，后面会反复用到：**一次 iteration 就是从模型的第 0 层一路 forward 到第 L−1 层、走一整遍。**喂进去的内容可以是某个 prompt 的一块、可以是好几个 request 的 decode 步、或者两者混着 —— 不管是什么，一次 iteration 把它们送过 L 层一遍。）\n想象一下引擎里的几秒钟。几十个 request 正在同时跑：有些还在嚼 prompt、有些已经 decode 到第 50 个 token、有些再吐 1000 步就要结束、还有些刚刚进来。新 request 进来，旧 request 走人。scheduler 的工作就是在伺候好每一个 request 的同时，把 GPU 塞得尽可能满。\n由此引出两个问题：\nrequest 的到达时间和结束时间各不相同。 引擎要怎么在不让谁在生命的起点或终点卡住的前提下，把 GPU 塞满？ 每次 iteration 的开销能差出 1000 倍。 一次纯 decode 的 forward 跑几毫秒就完事；如果其中混了一个 100 k-token 的 prefill，就要好几秒。怎么让每次 iteration 的开销大致平稳，scheduler 才好规划？ 两个答案 —— iteration-level 调度（ORCA）和 chunked prefill —— 用的是同一条直觉：我们要抚平的是每次 iteration 的开销，而不是每个 request 的开销。 ORCA 收拾\u0026quot;到达\u0026quot;和\u0026quot;结束\u0026quot;这两端的边界；chunked prefill 接着给一次 iteration 能装多少东西封顶。\n1. 朴素 batching 在哪些地方崩 想得到的最简单的 scheduler：有空位时挑出 B 个 request，一起过 prefill 和 decode，等最后一个跑完再把所有人的输出还回去，然后再挑下一批。这就是 request-level batching —— batch 是调度单位，batch 里有谁，在被收进来的那一刻就定死了。\n它撞上了真实流量的两条铁律。\nrequest 的到达时间不一样。 一个 batch 跑了 5 秒，第 200 毫秒进来的新 request 没法加进去 —— batch 里有谁，是在它开跑那一刻定死的。新 request 只能在队里等这一整批跑完。GPU 也许完全装得下再多一个 decoder，但引擎就是不放人进来。新 request 的 TTFT（time-to-first-token —— 你按下回车到 ChatGPT 的第一个字蹦出来之间的那段空白）从几毫秒膨胀到几秒，纯粹是干等。这就是 convoy effect：新到的人，就排在当前那批人里最慢那个的后面。\nrequest 的结束时间不一样。 一个 batch 里 request A 想要 50 个 output token，request B 想要 1000 个。两个一起 decode。第 50 步之后 A 已经做完了 —— 但它的 slot 收不回来给别人，因为 batch 的形状被钉死了，要等所有人都跑完才能动。接下来 B 还要再 decode ~5 秒，A 的算力 slot 就一直空着。更糟的是，A 已经生成好的那些 token，也得卡到 batch 边界才能还给用户。这就是 frozen batch size 问题：寿命短的 request 把最长那个邻居的寿命付了两遍 —— 一遍是返回延迟，一遍是 GPU 空转。\n两个失败都来自同一个根因：一个静态 batch 只有一个共享寿命，由它所有成员里最长的那个决定。 比这个最大值短的人在浪费；在开跑之后到达的人在干等。\nscheduler 被绑在了错误的粒度上。现实是按 iteration 这个粒度在动 —— 每次 forward，每个正在跑的 request 都产出一个 token（或者一块 prefill）。但 scheduler 在按 batch 这个粒度做决策 —— 几千次 iteration 才决策一次。当然跟不上。\n2. ORCA：按 iteration 调度，不按 batch ORCA 那篇论文 给的修法说起来很简单，后果却很大：把 iteration —— 端到端走完 L 层的一次 forward —— 当成调度单位。正在跑的 request 集合不再是收进来时定死的名单，而是 scheduler 在每次 forward 之间都在动手整理的一个活东西。\n两次 iteration 之间，scheduler 可以：\n踢人。 上一次 iteration 出 EOS 的任何 request，slot 立刻释放。 进人。 从队列里挑一个新 request 进来。它在第一次 iteration 就把 prompt 的若干行喂进去做 prefill。 带人。 还在 decode 中间的 request 接着跑，每个这一轮贡献正好 1 行 Q。 这三件事全是 scheduler 在 host 上做的 metadata bookkeeping —— 不动 GPU，只更新每个 request 的元数据。它们在两次 iteration 之间、GPU 还在忙上一次 forward 的同时跑完。\n这意味着一次 iteration 里都装了什么 比 Article 04 灵活了一档。Article 04 的 iteration 是同质的 —— 所有 request 都在 prefill，每个都贡献 prompt 的几行。ORCA 之下，一次 iteration 同时承载处于不同阶段的 request。举个具体例子：\nRequest 状态 这次 iter 的 Q 行数 kv_length A 第一次 iter 的 prefill，4096-token prompt 4096 4096 B decode 中间，第 51 步 1 1500 C decode 中间，第 200 步 1 1700 这次 iteration 的 Q 总行数：4096 + 1 + 1 = 4098。varlen kernel 沿这个 flat tensor 一个 request 一个 request 地走过去，算出三个互相独立的 score block：A 的 4096 × 4096（下三角 —— A 在给自己的 token 做 prefill）、B 的 1 × 1500、C 的 1 × 1700。每个 request 只读自己的 KV cache —— request 之间不会互相串。\n为了支持这种混合方式，Article 04 那个 cu_seqlens（只记 Q 行边界）泛化成每个 request 一个三元组：\n(q_start, q_end, kv_length) per request q_rows = q_end - q_start 是这个 request 这一轮贡献的 Q 行数。kv_length 是这次 iteration 把新的 K、V 追加进去之后，这个 request 完整的 attention 上下文 —— 也就是包含了之前 cache 里的所有内容。Q 行数和 kv_length 不再被强制相等 —— 一个 decoder 是 q_rows = 1 配 kv_length = 1500，一个全新的 prefill 是两者都等于 4096。\nkernel 这边就这一个改动。ORCA 真正的贡献不是一个新的 attention kernel —— 而是一种调度纪律：不要把一个 batch 跑到底，每次 iteration 都重新决定谁在里面。kernel 的活在 Article 04 里就已经准备好了；之前缺的是\u0026quot;按 iteration 一次次去用它\u0026quot;这条策略。\n现代 serving 系统说的 continuous batching，指的就是这件事。\n拿到的好处：\nConvoy effect 没了 —— 新来的人下一次 iteration 就能进来；等的是一次 iteration（几毫秒），不是一个 batch（几秒）。 Frozen batch size 没了 —— iteration t 释放出来的 slot，iteration t+1 就能填上；某个 request 跑完，EOS 一采样到就把结果还给用户，不必等到很远的 batch 边界。 两个问题都没了。漂亮的一仗 —— 漂亮到让人想就此宣布\u0026quot;调度问题搞定\u0026quot;然后翻篇。但我们偷偷跳过了一件事。\n3. 下一个问题：iteration 自身的开销也摇摆得厉害 ORCA 把\u0026quot;到达\u0026quot;和\u0026quot;结束\u0026quot;这两端的边界问题修了，靠的是把 iteration 升格成调度单位。但把 iteration 当成调度单位，同时也意味着它变成了整台引擎的心跳。所有正在跑的 request —— 不管是 decoder 还是 prefiller —— 每次 iteration 都各自往前走一小步。所以如果 iteration t 用了 6 ms、iteration t+1 用了 8 秒，那么任何正在跑的 decoder 的两个相邻 token 之间，都隔了 8 秒。一次 iteration 的 wall time 不再只是 GPU 内部花了多少算力的私事；它是这一轮里引擎里所有人共同的延迟下限。\n那 iteration 的 wall time 到底能波动多大？我们以 Llama-2-7B 在单张 H100 上跑为锚，把几种现实里 scheduler 真会拼出来的 iteration mix 套进成本模型走一遍。\n下面用到的 FLOP 和 wall-time 公式（点开看） Llama-2-7B：multi-head attention，32 层，hidden 4096，head dim 128。Forward 的开销有两个结构上不同的项。\nLinears：每行 token 过一遍模型的 forward 开销大致是 2P FLOPs，P ≈ 7×10⁹ —— 也就是每行 token 大约 14 GFLOPs。 Attention：每对 (q, k) 在整个网络里的代价 ≈ 4 · d_head · heads · layers = 4·128·32·32 ≈ 每对 0.52 MFLOPs。一段长度 L 的 prefill 配 causal mask：约 L²/2 对 ≈ 2.6×10⁵ · L² FLOPs。一个 decode step 对一份大小为 M 的 cache：M 对 ≈ 0.52·M MFLOPs。 H100 有效算力：fp16 compute-bound 任务 ~500 TFLOPs/s，read-bound 任务 ~3.35 TB/s HBM 带宽。一次 decode step 的主要开销是把整个网络的 weight 读一遍（fp16 下 ~14 GB），不是 FLOPs 本身 —— 所以 decode 是 bandwidth-bound 的，每步约 5–7 ms，由 bytes ÷ HBM 决定。\n底子是 8 个正在跑的 decoder，每个 context ~1 k。在同一次 iteration 里再塞进一个新 request，看几种情况：\nIteration mix（8 decode + …） Linear Attn Total Wall time 什么也不加（纯 decode） ~110 GF ~4 GF ~115 GF ~6 ms（被 weight 带宽卡住） + 1 k-token prefill ~14 TF ~0.3 TF ~14 TF ~30 ms + 4 k-token prefill ~57 TF ~4 TF ~61 TF ~120 ms + 16 k-token prefill ~225 TF ~67 TF ~290 TF ~580 ms + 100 k-token prefill ~1.4 PF ~2.6 PF ~4 PF ~8 s 三个值得注意的规律：\nLinears 跟 iteration 的总 token 数线性增长。 Attention 跟单个 request 的 prefill 长度平方增长 —— 短的时候可以忽略，到 100 k 这个量级开始压倒一切。 scheduler 合理拼出来的 iteration，wall time 之间能差到 大约 1300 倍。 正是最后这个数字把引擎搞崩。要感受一下它的含义：想象你正在用 ChatGPT，下一段正以每秒 ~150 token 顺顺地流着，突然 —— 在你能看见的范围里没有任何理由 —— 模型在一个写到一半的单词上冻住了八秒钟才接着写。你的对话本身一点没变。在你看不到的某个地方发生的事是：另一个用户往他自己的会话里粘了一份 10 万 token 的文档，你这一步 decode iteration 刚好被拼进了和他那个 prefill 同一次 forward。对 ORCA 来说这次 iteration 拼起来很正当 —— 两边都是合法的工作 —— 但 wall time 由他那个 prefill 决定，账由你来付。\n这种 head-of-line blocking 有两种表现形式，两种都发生在一次 iteration 之内，不是跨 batch。\n3.1 正在跑的 decoder 的 TBT 尖峰 上面那个场景有个名字：TBT（time-between-tokens） —— 一个 decode request 相邻两个 output token 之间的等待时间，决定了用户感受到的\u0026quot;匀速流式\u0026quot;那种体验。一个被 100k prefill 拖累的 iteration，会让所有恰好跟它共一轮的正在跑的 decoder 的 TBT 暴增 ~1300 倍。\n静态 batch 不会出这种问题 —— 但静态 batch 有它自己的灾难。ORCA 并没有破坏什么；它只是让一种本来就存在的差异在 iteration 这个层级上浮出水面，结果就是：它一出现，就同时砸在引擎里每个人头上。\n3.2 短 prefill 跟长 prefill 一起跑出来的 TTFT 尖峰 两个新 request 同一轮 iteration 一起到了：一个 prompt 100 token，一个 prompt 10 k token。ORCA 高高兴兴把它俩一起塞进同一次 forward —— 都是要做 prefill、都没有跑到一半的状态要照顾，往一次 iteration 里多塞东西本来就是 kernel 设计来做的事。但这次 forward 的 wall time 由长的那个邻居定：\nForward 内容 Linear Attn Wall time 单跑 100-token prefill ~1.4 GF ~3 MF ~4 ms 单跑 10 k-token prefill ~140 TF ~26 TF ~320 ms 100 + 10 k 一起跑 ~141 TF ~26 TF ~330 ms 短 request 的 TTFT 从单跑时的 ~4 ms 退化到了和长邻居一起跑的 ~330 ms —— 差了 ~80 倍，纯粹因为它俩共了一次 forward。从短 request 的角度看：整个网络对所有人都在全速跑，除了它；这个延迟在它自己的请求里找不到任何理由可以解释。这是结构性的 —— iteration 的 wall time 被里面最大那块决定，剩下的人就要陪着付。\n3.3 同一个根因 3.1 和 3.2 都来自同一个结构性事实：一次 iteration 的 wall time 由它里面最大那一块工作决定。 ORCA 可以决定一块工作要不要进这次 iteration，但决定不了一块工作有多大。在最大那一块被封顶之前，iteration 这个心跳就会跳乱。\n要把心跳搞稳，就得给最大那一块封顶。这正是 chunked prefill 做的事 —— 而 KV cache 已经替我们把工具准备好了。\n4. Chunked prefill：给最大那块封顶 如果长 prefill 是问题，那为什么不干脆把它切开？\n结构上其实没什么不让切的理由 —— KV cache 让\u0026quot;切开\u0026quot;这件事变得 trivial。chunk 0 跑完之后，它每一层的 K、V 已经存在 cache 里了。chunk 1 的 attention 直接读就行，跟 decode step 读 cache 是同一个动作。整段 prefill 一次性跑出来的数学和这个一模一样，结构上就是相等的；唯一的差别是这些活在什么时候被做。\n所以：把一段长 prompt 切成大小为 C 的 chunk，每次 iteration 带一块走。走一遍：一个长度为 N 的 prompt 在做 prefill 时会发生什么 ——\n这个 prompt 变成 ⌈N/C⌉ 次 iteration。 iteration 0 prefill token [0, C)。它的 attention 就是 Article 04 那种纯 prefill —— [C × C] 的下三角 score block。每一层的 K、V 存进 cache。 iteration 1 prefill token [C, 2C)。它的 attention 现在 Q 行来自这个新 chunk，K、V 行来自两边：cache 里的前缀，加上这一块刚算出来的 K、V。score block：[C × 2C]。 …… iteration k prefill token [kC, (k+1)C)。score block：[C × (k+1)C]。 chunk k 上的 mask 有两块：\n对着 cache 里前缀 的那块 [C × kC] 全开 —— 不 mask。前缀里的每个 token 都比这块里任何 token 早，causality 允许全看。 对着这块自己 token 的那块 [C × C] 是下三角 —— chunk 内部要保持 causal。 chunk k of size C, prefix S = kC: scores [C × (S+C)] cached prefix keys (S) this chunk's keys (C) queries from chunk k (C rows) attends to cached prefix all visible — no mask [C × S] causal [C × C], lower-tri masked q_rows = C, kv_length = S + C prefix block fully unmasked; new-token block lower-triangular 在一个大小为 C、前缀长度 S = kC 的 chunk 上走一遍 block：\n步骤 它做什么 碰 cache 吗？ LayerNorm 按行做 不碰 QKV proj 在 [C × hidden] 上 matmul → Q、K、V 各 [C × heads × d_head] 不碰 把 K、V 追加到 cache 把这一层的 K、V concat 到这个 request 的 cache 里 碰（写入） Attention Q [C × heads × d_head]、K 和 V [(S+C) × heads × d_head]。Scores [C × (S+C)]，mask 如上。 碰（读完整前缀） Output proj matmul 不碰 Residual + LayerNorm + FFN-up + GeLU + FFN-down + Residual 按行做 不碰 相比 Article 04 唯一变化的是 attention，几个形状被泛化了：\nQ 行数是 C，不再是\u0026quot;这个 request 的整个长度\u0026quot;。 K、V 行数是 S + C，不再等于 Q —— 前缀现在住在 cache 里。 score block 是矩形 [C × (S+C)]，不再是方阵。 Linears、residuals、layernorms、按元素的操作都是按行做的，根本看不见 cache。它们在 [C × hidden] 上一行一行处理，跟任何别的 C 行 batch 没有区别。\n这个 score block 的形状值得停下来想一想：它是一个混合体。左边那块 —— 这个 chunk 的 query 对前缀 cache —— 长得和把 C 个 decode step 摞在一起那种 score block 一模一样：对前面所有 token 全开 attention。右边那块 —— chunk 对自己 —— 是一个普通的 causal prefill 的 [C × C] block。decode 和 prefill 是同一个形状的两个极端，chunked prefill 是这个谱系上的任何一点。\n事后再看就一目了然：decode 不过是 C = 1 的 chunked prefill。 同一套机器，不同的旋钮取值而已。\n5. Piggyback：prefill chunk 跟 decode 共用一次 iteration 现在前面所有的零件要拼起来了。Article 04 那套 flat-tensor + varlen kernel 不在乎一个 request 的切片是哪种工作。对 kernel 来说，一个 request 的切片就是 (q_rows, kv_length) —— 不管它是在 decode（q_rows = 1）、是在 prefill 自己的第一个 chunk（q_rows = C, kv_length = C），还是在跑中间某个 chunk（q_rows = C, kv_length = S + C），形状都是这一个。\n所以一次 iteration 可以承载下面这些东西，全部打包进一个 flat tensor：\nIteration content: - Request E: prefill chunk 7 of 50 → 1024 Q rows, kv_length = 8 × 1024 = 8192 - Request A: decode step 51 → 1 Q row, kv_length = 1500 - Request B: decode step 200 → 1 Q row, kv_length = 1700 - Request C: decode step 75 → 1 Q row, kv_length = 1100 Total Q rows in this iteration: 1024 + 3 = 1027 varlen kernel 把每个 request 的切片各自走一遍。TP 还是完全没动到。\n这就是 piggyback chunked prefill：长 prefill 和正在跑的 decode 在一次 forward 里共存。scheduler 的工作变成了一种 bin-packing —— 每次 iteration 有一个预算（比如\u0026quot;Q 不超过 2048 行、iteration 不超过 50 ms\u0026quot;），用 decode step 和 prefill chunk 的任意组合把它填满。一段长 prompt 变成一连串 chunk 大小的贡献，每次 iteration 一块，和当时在跑的所有 decode 一起跑。短 prefill 一次就能跑完。decode 总是塞得下。§3 里那 1300 倍的波动塌缩成一个稳定的 iteration profile —— 大概 2 到 3 倍 —— 容易规划，引擎的心跳又稳了。\nC 是 scheduler 这边新增的旋钮：\nC 小 → iteration 时间更均匀、正在跑的 decoder 的 TBT 更低；但每个 chunk 的 cache 重读更多，linears 的 MFU 更低（小 GEMM 离峰值更远）。 C 大 → cache 重读更少、MFU 更高；但 iteration 的 wall time 又开始往上爬，所有人的 TBT 又开始劣化。 真实系统挑 C 一般落在 256–8192 这个范围，通常跟\u0026quot;每次 iteration 最多多少 token-row\u0026quot;的预算挂钩，预算的目标是控住 TBT 上限。举个具体的：在\u0026quot;每次 iteration ≤ 50 ms、最多 2048 个 Q 行\u0026quot;这套预算下，一段 100 k-token 的 prompt 会被 prefill 成 100 000 / 2048 ≈ 49 次 iteration，每一次都跟当时手上的 decode 一起跑。\n6. 成本分析 下面三件事值得停下来想一想，每一条都不是小开销。\n总算力没省。 chunk k = 0 … N/C − 1 各自的 C · (k+1)C 个 causal pair 加起来等于 N²/2。chunked prefill 把 attention 的活在 iteration 之间重新分配，但没有把它减少。\nKV 这边的 HBM 带宽消耗涨了。 chunk k 在每一层每一次 attention 都要再读 kC 行 cache。所有 chunk 加起来：≈ N²/(2C) 行 cache 累计流量，而不切的 prefill 只用 ~N 行（tiled attention 把 cache 流过去恰好一遍）。N = 100 k、C = 2048 时，同一段 prompt 的累计 cache 读带宽大约多了 25 倍 —— 这是 chunking 为了\u0026quot;把 iteration 封顶\u0026quot;付的价。这也是为什么 C 不能无限做小：到某个点之后，带宽税会盖过可调度性带来的收益。\nC 小的时候每次 iteration 的 MFU 会掉。 C 小的 iteration，linears 的 matmul 跑得离峰值远 —— tensor core 能咬的行数太少。真实 serving 引擎调 C 时找的是一个平衡点：让 iteration 时间能压到 TBT 目标、又不至于把 MFU 浪费太多。\n这三条加起来解释了 C 一般落在 256–8192 这个区间的原因。没有标准答案；区间具体落在哪儿，取决于模型本身的算力/带宽特性、以及引擎对 TBT 和吞吐的目标。\n7. 之后还有哪些新问题 到这里，我们手里有了一个真正的 serving 循环：prefill、decode、混合 iteration、单次 iteration 的开销有上限、没有空转的 slot。还有几条假设是漏的，每一条都给后面的文章埋下种子。\nKV cache 的物理布局。 我们一直默默假设每个 request 的 cache 在每一层都是一段连续的内存。一旦 B 涨起来、上下文长度又千差万别，这件事很快就难看 —— 碎片、eviction、分配开销。PagedAttention 把 cache 当虚拟内存来处理；下一篇文章。 两种 regime 共用一台引擎。 Decode 卡在 weight 读带宽上；prefill chunk 又是 compute-bound。也许它俩根本就不该共用同一组 GPU。Prefill/decode 拆分 探索的就是把它们丢到不同 replica 上去跑。 head 不总是独立的。 GQA、MLA 以及\u0026quot;减少 KV head\u0026quot;那一家子的其他成员，会大幅压缩 cache —— batch 更大、上下文更长 —— 但也引入一些此前我们可以无视的共享模式。一个子系列。 一个 request 的 cache 都装不下一张 GPU。 上下文长到 KV cache 自己就装不下一张卡的时候，sequence/context parallelism 会把一个 request 切到多张 GPU 上。一篇专门的文章，挪到后面。 每次都用同一种语法：松开一条假设，看看会掉出什么。\n","permalink":"https://wgzesg.github.io/llm_stories/zh/posts/05-orca-and-chunked-prefill/","summary":"很多 request 同时在跑，结束时间各不相同，有的还带着比一次 decode 大 1000 倍的 prefill。每次 iteration 的开销因此摇摆得厉害。ORCA 那种 iteration-level 调度先收拾一半问题；chunked prefill 再给最大的那次 iteration 封顶，让短任务不被拖在长任务后面。","title":"ORCA 和 chunked prefill：把每次 iteration 的开销摆平"},{"content":"Article 05 收尾的时候引擎跳得很稳。ORCA 把\u0026quot;进出 batch\u0026quot;的边界问题修了；chunked prefill 给每次 iteration 的开销封了顶，长 prompt 没法一个人霸占整间屋子。每次 iteration 都有上限、每个 request 都大致公平，引擎呼吸均匀。\n但那一篇文章末尾留了一根线，§7 第二条：\nDecode 卡在 weight 读带宽上；prefill chunk 又是 compute-bound。也许它俩根本就不该共用同一组 GPU。\n这一篇就是顺着这根线往下扯。我们用 roofline 的角度量一下这条沟到底有多宽，看着它在 context length 增长时越拉越开，最后落到结构性的修法：让两个阶段彻底不共用一台机器。\n起点其实有点尴尬：article 05 那套 piggyback chunked prefill 不是 prefill/decode 不匹配的答案，是个妥协。它把心跳抚平了，但底下的事实是 —— 一块 prefill chunk 和一个 decode token 对 GPU 的要求落在完全不同的 regime。共用同一次 forward，只是逼着双方都退到对自己不合适的那种 regime 里。\n1. Roofline，一页讲完 每张 GPU 上的每个 kernel，瓶颈都只在两种物理资源之一：\n算力（compute） —— tensor core 的峰值 FLOPs/s。 内存带宽（memory bandwidth） —— HBM 把 bytes 送到 SM 的速率。 （这里讲的是 GPU 内部的带宽，是 HBM 到 tensor core 之间那条管道。GPU 之间的带宽 —— NVLink、InfiniBand —— 是另一条轴线，等 TP/PP 出场我们再讲。）\n权重和数据实际放在哪：一张图 光说\u0026quot;内存带宽\u0026quot;很抽象。现代 GPU 有一套 memory hierarchy —— 几层缓存，越往上越小、越快。Tensor core 只能对 register 里的数据做运算，所以每一个 weight 字节、每一个 KV cache 字节，在真正被算到之前都得先沿着这条 hierarchy 走上来。\nGPU memory hierarchy（H100 风格的数字） GPU die SM 0 registers SRAM (~256 KB) tensor cores ~30 TB/s 有效 SM 1 registers SRAM (~256 KB) tensor cores SM 131 registers SRAM (~256 KB) tensor cores ⋯ ⋯ L2 cache ~50 MB 共享 ~5 TB/s 3.35 TB/s HBM 带宽 HBM — 80 GB 模型权重 · KV cache · 跨 kernel 的 activation 够装下模型和这一批的状态，但层级里最慢的一档 数字是 H100 风格的；其他 GPU 绝对值不一样，但形状 —— 顶层和底层在容量和速度上差三到四个数量级 —— 是普遍的。\n什么东西放在哪：\nHBM 装常驻的东西：模型权重（Llama-2-7B 是 14 GB）、每个 request 的 KV cache、跨 kernel 留存的 activation。容量大，相对慢。 L2 cache 是几个 SM 共享的一块小 scratch —— 多个 SM 都在读重叠数据时有用，但只有 ~50 MB，远装不下权重或 KV。 SRAM（per-SM shared memory） 是 kernel 当下正在动的那一块 weight、Q、K 的暂存区。FlashAttention 那套花招的核心就是把 attention score 矩阵压在 SRAM 里，不让它溢到 HBM。 Register 是 tensor core 真正读 operand 的地方。每个 SM 几百 KB，访问只要一个 cycle。 所以当你看到\u0026quot;kernel 从 HBM 加载了 14 GB 权重\u0026quot;的时候，路径是 HBM → L2 → SRAM → register → tensor core。一层比一层小、一层比一层快。3.35 TB/s 是这条链子最底的那一档 —— 也是一次 transformer iteration 没法绕过的那个瓶颈，因为权重比 HBM 之上每一层都大。\ncompute-bound 和 bandwidth-bound 在物理上到底是什么 矩阵乘按 tile 工作：从 HBM 把一块 A、一块 B 装到 SRAM 里，在 register 里相乘（每个元素背后是很多 FLOPs）、累加，下一块。同一块 weight tile 在被换出之前，会被多个输出行反复用到。\ncompute-bound 是 tensor core 跑满的状态。当前 tile 消耗得足够快，HBM 把下一块送来都来得及。带宽有富余。每个 weight 字节加载一次，被很多次 FLOPs 复用。 bandwidth-bound 是 HBM 送不上下一块。tensor core 已经把当前的吃完了，干等着字节到。每个字节复用的 FLOPs 太少，分摊不掉这一次加载的开销。 判断你处在哪种 regime 的那个数，恰好就是 每从 HBM 拉出一个 byte，做了多少 FLOPs —— 这就是 intensity。也是为什么 roofline 这条规则没什么回旋余地：它不是经验观察，是上面这套 hierarchy 的直接推论。\nRoofline 这条规则 哪种资源是瓶颈，由一个数决定：arithmetic intensity I，FLOPs 数和从 HBM 加载的 bytes 数的比值：\nI = FLOPs done / bytes loaded （单位：FLOPs/byte） 硬件这边对应有一个数，叫 ridge point R：\nR = peak FLOPs/s / peak HBM bandwidth （单位：FLOPs/byte） H100 SXM5：fp16 GEMM 持续算力 ~500 TFLOPs/s，HBM3 带宽 3.35 TB/s → R ≈ 150 FLOPs/byte。\n规则：\nI \u0026gt; R → compute-bound。算力是瓶颈；带宽有富余。 I \u0026lt; R → bandwidth-bound。字节是瓶颈；tensor core 在等数据。 就这一条。剩下整篇文章其实就是在反复问两个问题：\n一次 prefill iteration 和一次 decode iteration 的 I 各是多少？ context length 增长时 I 怎么变？ 2. 纸面估算单次 iteration 的开销 先把符号定下来。全文假设 fp16（每个参数 2 byte、cache 里每个数也 2 byte）。换成更低精度的 dtype，数会跟着变，但故事不变。\n符号 含义 单位 Π 参数总数 无量纲 K_tok 一个 context token 在 KV cache 里的字节数（所有层的 K + V 加起来） bytes/token T 这次 iteration 里的 token 总数 tokens B 这次 iteration 里在跑的 request 数 无量纲 L 每个 request 的平均 context length tokens C prefill chunk 大小（每个 chunk 的新 token 数） tokens R 硬件 ridge point FLOPs/byte （K_tok 是把所有层加起来的总和 —— 一个 context token 在整张网络的 KV cache 里要占多少字节，不是每层。）\n先点一下：transformer block 里大部分算力和几乎全部参数都在它的 matmul 层 —— QKV projection、attention 的 output projection、FFN 的 up/down projection。Attention 本身（softmax over scores 那一步）和 pointwise 操作（layernorm、GeLU、residual add）只占总 FLOPs 的一小块（除非上下文极长）。所以下面说\u0026quot;每个 token 多少 FLOPs\u0026quot;或者\u0026quot;weight bytes\u0026quot;的时候，意思都是 matmul —— 那才是开销所在。\n对一次跑在 Π 大小模型上的 iteration，有两个物理量要追，两个都关于 Π 线性：\n从 HBM 拉出来的 weight bytes： 每个参数 fp16 = 2 byte，一次 iteration 读一遍 → 2Π bytes。Llama-2-7B（Π = 7B）就是 14 GB。一次 iteration 付一次，跟塞了多少 token 没关系。 每个 token 走完整张网络做的 FLOPs： 一个 token 走过一层 matmul，和那层每个参数都做一次 multiply-accumulate（每个参数 2 FLOPs）。把整张网络的 matmul 加起来，每个 token 2Π FLOPs —— Llama-2-7B 就是 14 GFLOPs/token。一次 iteration 处理 T 个 token，做 2Π · T FLOPs。token 在 matmul 里互不干扰（只在 attention 里互相看到），所以同一次 iteration 里两个 token 的算力是单个 token 的两倍 —— 但 weight 加载只付一次。 加上 KV cache 读取，把两件事一起写下来：\nbytes_loaded = 2Π (weights, 一次 iteration 付一次) + K_tok · L · B (KV cache, 每个 request 读自己那 L 行) FLOPs_done = 2Π · T (T = 这次 iteration 里的 token 数) 塞进 intensity 的定义里：\nI = 2Π · T / (2Π + K_tok · L · B) 盯着这条公式看一会儿 —— 这一节余下的内容就是把它读仔细。分母两项、分子一项；按顺序走一遍，整个 prefill/decode 故事就出来了。\n第一步：先假装 KV 这一项是零 在 L 极短、或者一段对话刚开始还没什么 context 的时候，分母由 2Π 主导，公式塌成：\nI ≈ T intensity 就是 共享同一次 weight 加载的 token 数。prefill 和 decode 在这里就分了岔：\nPrefill iteration：T = C = 2048 个 token → I ≈ 2000 → 远高于现代 ridge point（~150） → compute-bound。 Decode iteration：T = B（同时在 decode 的 request 数，一般几十到一百多）→ I ≈ B → 远低于 ridge → bandwidth-bound。 同一张 GPU、同一个模型、同一个 kernel。唯一区别是这次 iteration 装了多少 token。Prefill 把一次 weight 加载分摊到几千个 token 上；decode 分摊到 B 个上。从第一次 iteration 起，它们就坐在 ridge 的两边 —— 而且差距不小：intensity 上至少差一个数量级。\n凭直觉想到的修法是把 decode 的 batch 推得更大 —— 把 B 推到 intensity 越过 ridge 为止。要清掉 R = 150，得 B ≥ 150。下一步说为什么这条路走不通。\n第二步：把 KV 那一项打开 context 一长，K_tok · L · B 就开始往分母里加。两项相等的位置（crossover）：\nL · B = 2Π / K_tok Llama-2-7B（Π = 7B、K_tok ≈ 512 KB）下，L · B ≈ 27 k。decode batch B = 32 的话，crossover 落在 L ≈ 850 token。\n850 这个数，放到今天的标准里小到吓人，值得停一下。生产环境里的 prompt 现在动不动就是几万 token：超长的 system prompt 和工具定义、RAG 灌进来的文档、累积的多轮对话、agentic chain 那种 input/output ratio 经常 100:1 起步的工作流。前沿模型出 200 k – 2 M 的 context window，是因为真实的 workload 真的会塞满。所以\u0026quot;过了 crossover\u0026quot;根本不是 corner case，而是中位 request。\n过了 crossover，公式向另一个方向化简：\nI ≈ 2Π · T / (K_tok · L · B) 这里的约分关系开始决定命运：\nDecode（T = B）：I ≈ 2Π / (K_tok · L)。B 上下消掉 —— *在长 context 下，把 decode 的 batch 推大不再能提升 intensity。*多收的 request 只是按比例多付 KV 读带宽。再加上 KV 内存预算，B 还涨不太大就先把卡撑爆。所以\u0026quot;把 batch 推大\u0026quot;这一招在第一步本来就走不通，到了第二步又会再撞一次墙。 Prefill（T = C）：I ≈ 2Π · C / (K_tok · L · B)。没东西约掉 —— C 老老实实留在分子里。Prefill 一直 compute-bound，到夸张的 context 长度都还守得住。 同一条公式，两件事 Prefill 是 compute-bound、decode 是 bandwidth-bound。 在 context 接近零的时候就成立，完全由\u0026quot;同一次 weight 加载分摊到几个 token\u0026quot;决定。两个阶段从一开始就坐在 ridge 的两边。 长 context 把这条沟拉得更宽。 第二项带宽成本（KV 读）从分母里冒出来，过了 crossover 就主导（生产流量基本都过 crossover）。受影响的主要是 decode 这边，prefill 几乎没事。 §3 用 Llama-2-7B 上的具体数字把这两件事坐实。\n3. 一个模型、两个阶段、两张表 把公式落到地面上：在一个具体模型 + 一张具体 GPU 上，扫一遍 L。\nLlama-2-7B（MHA、32 层、32 head、head_dim 128、fp16）on H100：\nweight bytes 2Π = 14 GB K_tok = 2 (K,V) · 32 层 · 32 head · 128 head_dim · 2 byte ≈ 512 KB/token ridge R ≈ 150 FLOPs/byte Decode at B = 32 L weight bytes KV bytes total I = 2Π·B / total regime 1 k 14 GB 16 GB 30 GB ~15 bandwidth-bound（weights ≈ KV） 4 k 14 GB 64 GB 78 GB ~5.7 bandwidth-bound（KV 主导） 16 k 14 GB 256 GB 270 GB ~1.7 严重 bandwidth-bound 64 k 14 GB 1.0 TB 1.0 TB ~0.4 cache 在一张 H100 上装不下 （分子 2Π · B = 448 GFLOPs —— 钉死的。是分母在炸。）\n注意几件事：\nIntensity 跌得很快。 从 L=1k 的 ~15 跌到 L=64k 的 ~0.4 —— 单一个 context 维度上就掉了一个数量级以上。 内存预算先于带宽爆。 L=16k、B=32 时单 KV 就 256 GB，远超 H100 的 80 GB。PagedAttention 之所以存在，一部分就是为了管这件事；B 在长 context 下被迫往下压，结果 intensity 又被进一步拖坏。（Llama-2-7B 用的是 MHA；现代 GQA/MLA 把 K_tok 砍掉 4–8 倍，主要就是为了把这堵墙往后推。） 主导的字节种类会变。 短 L 时 weight 主导，长 L 时 KV 主导。两边都是 bandwidth-bound，但解法不同 —— batch 推大对 weight 带宽有用；GQA/MLA/FlashDecoding 对 KV 带宽有用。 Prefill at C = 2048 Chunked prefill 拿 C 个新 token 跑，面对一段长度 S 的前缀（所以 T = C 个 token 的算力，读 S 个 token 的 cache KV）：\nI_prefill = 2Π · C / (2Π + K_tok · S) 分子里挂着 C —— 每个加载的字节都被分摊到几千个 token 的算力上。\n前缀 S weight bytes KV bytes total I regime 4 k 14 GB 2 GB 16 GB ~1800 compute-bound（高出 ridge ×12） 64 k 14 GB 32 GB 46 GB ~620 compute-bound（×4） 256 k 14 GB 128 GB 142 GB ~200 还是 compute-bound（×1.3） 1 M 14 GB 512 GB 526 GB ~55 终于跌到 ridge 之下 —— 但已经 100 万 token 了 Prefill 能一直 compute-bound 撑到极端的 context 长度。哪怕真的跌到 ridge 下面，也远远没到 decode 在常见 context 长度下那种 bandwidth-bound 的程度。\nasymmetry 一句话说清楚：\nPrefill 把每一字节带宽分摊到 C ≈ 2000 个 token 上；decode 是每个 request 一个 token。长 context 把刀子往 decode 这边拧，prefill 几乎没动。\n同一个模型，同一张 GPU。两个阶段。两条完全不同的命运曲线。\n4. 为什么一台引擎没法把两边都伺候好 把 article 05 那台引擎 —— continuous batching、chunked prefill、piggyback iteration —— 拿过来问一句：你怎么 sizing？\n按 prefill sizing： 给 GPU 选高 FLOPs 的型号。Decode 跑在一种 ~90% 算力天生用不上的硬件上，因为它是 bandwidth-bound。你为 decode 物理上用不到的 tensor core 付钱。 按 decode sizing： 选少一点、按 HBM 带宽和容量来挑的 GPU。Prefill 跑在缺 FLOPs 的机器上，时间被拖长。TTFT 上去。 混跑： 每次 iteration 把 prefill chunk 和 decode token 装在一起。TBT 被这次 iteration 里 prefill chunk 抢走的那点算力扣下当人质。Chunked prefill 给这条卡了上限 —— 那就是 article 05 的全部目的 —— 但这个上限不是免费的。共用一台引擎的一次 decode iteration，要为 C 行根本对自己没用的 prefill 算力买单。 更深一层：**workload 的瓶颈 profile 是双峰的，引擎是单峰的。**没有一种 sizing、没有一种 parallelism 策略、没有一种 batch policy，能同时把两个阶段都伺候好。两个阶段拉满的是不同的物理资源、追的是不同的 SLO（TTFT vs TBT），一个 scheduler、一个旋钮，没法在两种 regime 下同时满足两个 SLO。\n那就别拼了。建两个 pool。\n5. 拆开 prompt Prefill pool compute-bound 目标是 TTFT 无状态 KV cache transfer 每个 request, L_p · K_tok bytes Decode pool bandwidth-bound 目标是 TBT 持有长寿 KV tokens 一个 request 的生命周期中间多了一跳：\nPrefill pool 收下 prompt，对全部 L_p 个 token 跑 chunked prefill，产出整个 request 的 KV cache 加上第一个生成 token。 KV cache 传输 把这 L_p · K_tok byte 从 prefill GPU 内存搬到 decode GPU 内存。 Decode pool 接到 KV cache，把 request 塞进自己的 continuous-batching 池子，一直 decode 到 EOS，把 token 流式吐回给用户。 两个 pool、两套调度、两个 SLO 目标。妥协没了。每个 pool 现在可以针对一个单一目标自由地选 parallelism、batch policy、硬件搭配、scheduling 纪律。这份自由就是最大的那块收益 —— 各自具体怎么用它，留给系列后面的文章。\n新成本是中间这一跳。我们在 §6 给它定价。\n6. 新成本：KV cache 传输 拆开两台引擎之后，每个 request 都要把 KV cache 从一边往另一边搬一次。这是真成本，先估个数。\nLlama-2-7B（K_tok ≈ 512 KB）一段 4 k-token 的 prompt：\n每个 request 的 KV bytes = L_p · K_tok = 4096 · 512 KB ≈ 2 GB 每个 request 2 GB。每秒几百个 request（不算高的生产负载）的话，两个 pool 之间的总 east-west 流量轻松能到几百 GB/s。中间那条 fabric 得吃得下这个量。\nfabric 长什么样、一次传输要多久：\nFabric 带宽 传 2 GB 要多久 NVLink（节点内） ~900 GB/s ~2 ms NVLink-network / NVSwitch fabric（集群内） ~400 GB/s ~5 ms InfiniBand HDR（跨节点） ~50 GB/s ~40 ms PCIe Gen5（host 中转） ~64 GB/s ~30 ms 所以两个 pool 同在一个 NVLink domain 里的话，这一跳几乎免费；隔了 IB 的话，是真的税。40 ms 加在 TTFT 上能感觉到，5 ms 是无所谓。\n由此马上冒出几个工程旋钮（每一个都够独立成篇 —— 这里我们只点出来，不解决）：\nLayer-streaming overlap. 别等 prefill 全跑完再开始传。每一层的 K、V 是按顺序产出的；后面层还在算的时候，前面层的 KV 就已经能往那边发了。做得好的话，传输几乎完全藏在 prefill 算力背后。 GPUDirect RDMA. 字节直接在两块 GPU 的 HBM 之间走，不绕 CPU 内存。省掉一次拷贝、一次 context switch。 拓扑感知调度. 把同一个 request 的 prefill 和 decode 排到拓扑上靠近的 pool —— 同机架、同 NVLink domain —— 把 fabric 那一档压低。 前缀复用. 两个 request 共享一段长前缀的话，只需要算和传 suffix 那段的 KV。生产系统（Mooncake 是个写得比较细的例子）把这件事做成了内存层级的问题：热前缀在 HBM、温前缀在 DRAM、冷前缀在 SSD。 GQA / MLA 直接砍单价. 把 K_tok 砍掉 4–8 倍，传输也跟着砍 4–8 倍。一般不把它叫做拆机优化，但实际上是。 每一条底下都能再开一篇文章。这一节的 takeaway 就是：这次传输是拆机的代价，但是付得起的 —— 有上限、能工程化、相对于 TTFT/TBT 上的收益来说很小。\n用户感受：\nTTFT = prefill time + transfer time + 第一次 decode iteration。transfer 是真的在里面，但量级是几 ms 到几十 ms。 TBT = 纯 decode，不会被 prefill 抢算力。decode pool 的每次 iteration 都只装 decode 工作，所以 TBT 平稳到 decode 硬件能给到的极限。 这桩交易就是想要的那种：TTFT 上一次性吃个小亏，换取整段生成里 TBT 又稳又可预测。用户对 TBT 的感受比对 TTFT 重得多 —— TTFT 是一次顿挫，TBT 是每一次顿挫。\n7. 之后还有哪些新问题 Article 05 是给 iteration 封顶。Article 06 是把它拆开。§2 那条公式逼着我们答了\u0026quot;为什么要拆\u0026quot;；这一篇大部分篇幅都在做这件事。\u0026quot;怎么让它真跑起来\u0026quot;是另一个问题，大部分实际工程量都落在这一边 —— §6 应该读成一扇门、不是终点站 —— 它只是更大一片工程面里露出水面的那一截。\n在那扇门口站一会儿。两块 GPU，可能在不同机架、可能挂在不同的内存层级下，要在能藏在 prefill 延迟里的时间内，把以 GB 计的状态搬过去。这条管道里每一个选择背后都有一片很实在的设计空间：\n字节走哪条 fabric —— NVLink 还是 NVSwitch 还是 InfiniBand 还是 PCIe —— 单次传输的成本能差出近两个数量级（§6 的表格）。你建出来的集群拓扑长什么样，全看这道选择。 request 之间 KV cache 住在哪里 —— HBM 还是 DRAM 还是 SSD —— 让拆机引擎变成了一套分层的内存系统。Mooncake 那套前缀池是其中一种；还有别的实现，invalidation 和 locality 行为各不相同。 传输怎么和算力 overlap —— layer-by-layer streaming、GPUDirect RDMA、双缓冲队列 —— 这些是让一跳\u0026quot;端到端看不见\u0026quot;还是\u0026quot;在 TTFT 里非常显眼\u0026quot;的分水岭。 request 怎么在两边 pool 之间路由 —— fabric 局部性感知调度、前缀缓存命中、decode 容量追踪 —— 在 article 05 的所有调度问题之上又叠了一层。 每一条都能独立成篇，系列下一篇要接的就是这根线 —— 跑一套拆机 serving 系统的工程问题。然后 我们才能干净地问下一组问题：拆机让我们终于能干净地问的那些优化问题。每个 pool 现在有了专门化的自由，它想要什么？Pipeline parallelism 给 prefill、Tensor parallelism 给 decode、PagedAttention、GQA/MLA、FlashDecoding、speculative decoding —— 每一项在 pool 拆开之后都有了一个干净的位置，我们再按顺序一个一个走过。\n每次都是同一种语法：找出瓶颈、把 workload 切到每块只看见绑住自己那一种瓶颈的程度、按块去优化。拆机是这种切法里最大的一刀。系列接下来要做的，是这一刀切出来之后该做的工程和优化。\n","permalink":"https://wgzesg.github.io/llm_stories/zh/posts/06-prefill-decode-disaggregation/","summary":"Article 05 让两个阶段勉强共用一台引擎。这一篇要说的是：它俩本来就不该共用 —— prefill 是 compute-bound、decode 是 bandwidth-bound，长上下文还把这条沟越拉越宽。承认了这种 asymmetry，拆机就不再是优化，而是顺着公式来唯一说得通的答案。","title":"Prefill/Decode 拆机：两个阶段坐在 roofline 的两边"}]