VLLM on Fain的Blog

Pico-vLLM 开发日志 #14 Prefix Caching

Sat, 18 Apr 2026 21:56:39 +0800

“如果你想体验极致的数学规整美感和爽感，就去实现Paged Attention；如果你想体验抽象分层设计原则和一致性维持的究极长痛，去实现Prefix Caching。”

——我自己

======================================================================

这么多天没更新，今天终于更新了。之所以没有和之前一样2~3天一更新，是因为这中间的间隔，几乎全部在对Prefix Caching进行架构设计、实现迭代和反复的DEBUG。个人的体感是，在所有的推理框架实现当中，Prefix Caching绝对是除了分布式中最复杂的部分（当然，这部分非常非常大，但对于许多feature来说，其实并不是和分布式本身强绑定的）以外的皇冠上的明珠，是实现最为复杂的、同时也是相当重要的架构级feature，其挑战在于其并非是局部的算子或者数据结构设计，而是同时牵动存储分配、请求生命周期、缓存一致性的全局性架构。尤其是，如果需求足够好的兼容性，也就是需求既有radix-tree的token级别的粒度和相应可扩展性，又需求底层的Block Manager能够天然、不需要额外处理的自然同时兼容radix-tree模式，会引入非常复杂的系统架构设计需求和要求非常小心的实现。因此，这里会完整的把我的开发工程、主要最后定稿的想法和设计模式、途中遇到的一些问题，尽可能在遗忘和丢失之前。如果能够帮助到后来者和读者的话，那就太好了。

架构设计：在 vLLM 的地基上建 SGLang 的房子

LLM 推理的一个核心优化是前缀复用，这已经被多篇论文证明在特定的工作场景下非常有效。它的思想很简单：多个请求共享相同的 system prompt 时，只需计算一次 KV Cache，后续请求直接复用，通过指向相同的KV Cache，实现在无复制、无额外空间占用的情况下完成相同功能的目的。从逻辑上来说，它是对Paged Attention和分页、页表思想的自然延伸。从工程上来说。SGLang和vLLM分别用两种截然不同的方式实现了这个优化。下面的描述中就可以看出，vLLM的架构设计模式更多是在naive Paged Attention上的一个优美的最小侵入改动以部分的模拟前缀匹配，而SGLang则就真的是毫无妥协的、原生的基于完整一体化设计的Prefix Caching。

在我自己进行实现的时候，一开始其实并没有意识到一个重要问题：两者的差距其实并不小，而且其实是基于各自的先决条件（Block Size的不同默认），也没有意识到为什么vLLM要使用哈希表这种不足够优美的讨巧的“模拟radix tree”的路径。我在刚开始这个feature实现的想法是：vLLM应该是出于性能原因，或者架构太难以改动，因此做出了设计上的妥协，我应该完全能够在SGLang风格下对我自己的框架进行修改。当中途我发现这个过程里面的坑多的令人头皮发麻的时候，代码已经改的太多，无路可退了。因此，我本人几乎是非常头铁的一路撞到了底，硬生生把东西全都做了出来。不过，如果让我自己再做一遍的话，我可能应该会直接采用vLLM模式：除了可扩展性略差之外，在性能上并没有实质性的差距，而且实现要简单的多。

vLLM模式

vLLM的设计采用的是block粒度的hash table。以block_size（通常16）个token为一组，计算hash后存入字典。匹配只发生在block边界上。如果两个prompt在block中间分叉（只要block的id是不一样的），整个block都被计为miss，不会进入前缀匹配的范围。

SGLang模式

SGLang用token粒度的radix tree。每个token都是树的一个路由单元，KV Cache的存储粒度（page_size）也是1个token。树和存储粒度一致，匹配可以精确到任意token位置，不存在对齐问题。

粒度不一致的桥接

Pico-vLLM的底层存储是vLLM 风格的BlockManager（block_size=16），配合手写的Triton kernel实现PagedAttention。但前面也提到，我想要的SGLang风格的匹配模式。也就是说，token级别的radix tree，任意位置分叉都能在radix tree级别正确处理。

问题是：token粒度的索引和block粒度的存储，天然不兼容。这个问题的处理占据了我大量的开发时间，也引入了相当多的bug。不过，我最终得以找到一个相对优美的解决方案，如下所示。

三层架构

最终的设计是一个三层结构，中间有一个显式的桥接层：

RadixTree（token 粒度）
 ↑ match_prefix: 在 token 级别遍历，返回 block-aligned 的结果
 ↑ insert: 在 token 级别分裂/合并，block_ids 按边界分配
 |
PrefixCache（桥接层）
 ↑ 向下取整对齐：matched_len = (raw_len // block_size) * block_size
 ↑ 双层引用计数协调
 |
BlockManager（block 粒度）
 ↑ allocate / free / inc_ref / dec_ref
 ↑ swap_in / swap_out（offload 预留）

RadixTree 的节点存储的 cached_blocks 是逻辑 block id 列表。一条边可能有 90 个 token 但只覆盖 5 个完整 block（80 token），剩余 10 个 token 的 KV 在第 6 个 block 的前 10 个 slot 里，但这个 block 不完整，不能被 prefix cache 复用。

match_prefix 在遍历完所有匹配的边之后，在返回前做一次统一对齐：

1
2
3


aligned_len = (matched_len // self.block_size) * self.block_size
aligned_blocks = matched_blocks[:aligned_len // self.block_size]
return aligned_blocks, aligned_len, last_node

这个设计把对齐逻辑集中在一个点上，无论树的内部结构怎么分裂合并，返回给外部的永远是 block-aligned 的结果。Engine 不需要知道 radix tree 内部的 token 粒度。如果后续实现写时复制，这个部分就可以进行改变，此刻以预留的形式存在。

insert 的 block 归属

insert 触发 split 时，需要把原节点的 cached_blocks 正确分配给 split_node 和 child。一条 96-token 的边在第 90 个 token 处分裂：

split 前：node [t0..t95] cached_blocks = [B0, B1, B2, B3, B4, B5] (6 blocks, 96 tokens)

split 后：
 split_node [t0..t89] cached_blocks = [B0, B1, B2, B3, B4] (5 blocks, 80 aligned tokens)
 child [t90..t95] cached_blocks = [B5] (1 block, 但只填了 6 token)

block 归属的计算：num_split_blocks = (i + match_len) // block_size - i // block_size。这个公式处理了分裂点不在 block 边界上的情况。前半段拿走所有完整 block，后半段继承剩余的。

语义设计：处理边界情况和分层解耦

（以下内容有ai辅助排版）

两层引用计数的必要性

SGLang 只有一层引用计数（lock_ref），因为它的 page 分配器和 radix tree 是一体的。vLLM 的 hash table 也只需要一层 block 引用。但 Pico-vLLM 的 radix tree 和 BlockManager 是独立的模块，而它们各自需要回答不同的问题，因其抽象层级和最细粒度同时不同。具体来说，其被设计的按如下职责工作：

BlockManager.logical_ref_count的含义是"这个block此刻有几个持有者"

持有者只有两类：正在运行的请求（通过 kv_cache 持有）和 RadixTree（作为一个整体实体持有）。ref归0意味着block立刻回收到 free pool。这是决定"block 是否可回收"的唯一依据。

RadixTreeNode.lock_ref的含义是"这个prefix段还有几个活跃请求在用"

lock_ref > 0意味着有请求正在使用这段前缀的 KV Cache。驱逐这个节点不会导致恶性的数据丢失（因为请求的持有阻止block被释放），但会引发额外的处理复杂性。在目前策略下，lock_ref归0意味着节点变得可以最小代价驱逐，进入LRU队列等待清理。

核心不变量：

logical_ref_count[block] =
 #{ 活跃请求 r : block ∈ r.kv_cache.logical_block_ids }
 + 𝟙[ 存在 radix 节点 n : block ∈ n.cached_blocks ]

这和文件系统的 inode + hard link 模型完全同构——BlockManager 是 inode table，RadixTree 是目录树。"block 什么时候可以物理回收"和"目录项什么时候可以删除"是两回事。RadixTree 无论内部有多少个节点引用了同一个 block（比如 split 前后的不同节点），在 BlockManager 看来都只算一个持有者。

lock_ref 的 node 指针方案

这是整个设计中最关键的决策，也是踩了最多坑之后才得出的结论。

最初的实现用 token 路径做 inc/dec：match 时沿 tokens 遍历 inc，close 时沿同样的 tokens 遍历 dec。看起来对称——但 insert 会改变树结构。

Bug 1：release 用 match 重新找 block

三个请求使用相同的 prompt。req4 首先 prefill，insert 把 blocks [A, B, C] 加入 radix tree。req5 在 submit 时 tree 为空（所有 submit 发生在 step 之前），matched_blocks = []，自己 allocate 了 [D, E, F]。insert 时发现路径已在，newly_held = []。

req5 close 时，release 调用 radix_tree.match(tokens) 重新找 block——找到了 req4 的 [A, B]，然后对不属于自己的 block 执行 dec_ref。A、B 的 ref 从 1 降到 0，被释放。req6 close 时又 match 到已经释放的 [A, B]，继续 dec，ref 变成 -1。

req5 close:
 match 返回 [A, B]（属于 req4 的 block）
 block_mgr.dec_ref([A, B]) → A=0 free, B=0 free ← 错误释放！

req7 submit:
 match 返回 [A, B]（tree 指向已释放的 block）
 adopt_blocks 读到 block_mapping[A] = (NONE, -1) → 崩溃

修复：release 不再 match，而是接收 held_blocks = kv_cache.logical_block_ids（请求自己实际持有的 block 列表）。

Bug 2：split 后边长不对齐 block_size

请求 A 的 prompt 有 144 token，insert 建了一条 144-token 的边。请求 B 的 prompt 前 125 个 token 相同，第 126 个开始不同。insert 触发 split：

split_node [t0..t124] len=125 ← 不是 16 的倍数！

请求 C 来了，match 走到 split_node，完整匹配 125 个 token。返回 matched_len=125。Engine 调用 adopt_blocks(blocks, 125)，assert 125 % 16 == 0 失败。

修复：match 返回前统一做 aligned_len = (matched_len // block_size) * block_size。

Bug 3：token 路径遍历的根本缺陷

即使修了前两个 bug，第三个更深的问题出现了。请求 R 在 submit 时 match 到 112 个 token（对齐后），inc_ref(tokens[:112]) 沿路径遍历。此时树有一条 144-token 的边，边长 > 112，inc_ref 走不到这条边，没有计数。

然后 R 的 insert 触发 split，把 144-token 的边劈成 split_node（125 token）+ child（19 token）。split_node.ref 复制了原节点的 ref。

R close 时，radix_held_len = 112 + ext * 16 = 144。dec_ref(tokens[:144]) 这次能走到 split_node（125 ≤ 144），对 split_node 执行 dec。但 inc 时没碰过 split_node——inc 了 0 次，dec 了 1 次，underflow。

inc_ref(tokens[:112]): 边长 144 > 112 → break，不 inc
 （insert 触发 split）
dec_ref(tokens[:144]): split_node 边长 125 ≤ 144 → dec → 0-1 = -1 💥

根本原因在于用token路径重新遍历做 inc/dec，但树的结构在match和close之间被insert改变了。 同样的token序列在两次遍历中走到了不同的节点。

修复方案：node指针方案

SGLang的lock_ref用的是node指针+parent chain，而不是token路径重遍历。最终采用了这个方案：

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


def match_prefix(self, tokens):
 ...
 last_node = self.root # 初始指向 root
 while ...:
 if match_len == len(edge):
 last_node = child # 只在完整匹配边时更新
 ...
 return matched_blocks, aligned_len, last_node

def inc_lock_ref(self, node):
 curr = node
 while curr is not None and curr is not self.root:
 curr.lock_ref += 1
 curr = curr.parent

def dec_lock_ref(self, node):
 curr = node
 while curr is not None and curr is not self.root:
 curr.lock_ref -= 1
 curr = curr.parent

请求对象上存的是 request.last_node（指针）而不是token路径。close时dec_lock_ref(request.last_node)通过parent chain向上遍历以正确处理引用数减少问题。

**这个设计方案的主要目的是为了确保引用技术操作在结构变化下永远对称。**split 只是在 node 和 root 之间插入了新的中间节点（split_node），Python 对象的 parent 指针会被正确更新。从 child 走到 root 的 parent chain 自动包含了新插入的 split_node，不需要重新遍历 token。

split 前：root → child [t0..t143] lock_ref=K
 inc_lock_ref(child): child.lock += 1

split 后：root → split_node [t0..t124] → child [t125..t143]
 split_node.lock = child.lock（纯复制）
 child.parent = split_node（指针更新）

dec_lock_ref(child): child.lock -= 1 → split_node.lock -= 1 ← 自动走过新节点，对称！

insert部分不改变lock_ref的计数。新节点创建时 lock_ref=0，分裂时 split_node.lock 纯复制自 child.lock。lock_ref 的增减完全由 match 和 close 控制，insert 是透明的。

权责清单

最终的修改入口和权责划分：

事件	BlockManager 引用计数动作	RadixTree 引用计数动作
allocate(n)	新 block +1	—
PrefixCache.match	matched_blocks +1	inc_lock_ref(last_node)
PrefixCache.insert	newly_held +1（RadixTree 新持有）	新节点 lock=0
PrefixCache.release	held_blocks -1	dec_lock_ref(last_node)
PrefixCache.try_evict	evicted_blocks -1（RadixTree 停止持有）	lock_ref=0节点从树删除

关键设计原则在于每一行的两列操作完全独立，互不依赖。BlockManager不知道RadixTree的存在，RadixTree不直接调BlockManager的free，PrefixCache是唯一协调两者的层。其中，加粗部分是可以根据不同的驱逐策略进行联合设计和改变的部分，只要确保可驱逐和操作成对合法，完全可以对这些部分进行改变，只要正确处理node废弃之后的状态、指针引用即可。

实际上，如果驱逐策略变得很复杂，两层引用计数的重要性将会开始凸显。这是因为实际上Radix Tree的引用计数本身不是释放是否合法的硬性指标（它是虚拟层级的最上层），而是“判断是否应该驱逐、应该驱逐哪些”的参照性标准。从技术本质上来说，在radix tree层面上驱逐一个计数不等于0的块是完全合法的，因为Block Manager此时并不会真的立刻释放它，而是给了Block Manager一个可以尽快在合法条件下按需将其释放的信号（或者，反过来说，radix tree是阻止block Manager将相应块释放的控制机制），这就赋予了极大的自由度（和处理被驱逐的废弃节点的相应活跃请求的复杂性）。对于需要考虑换入换出成本的radix tree的offload相关的驱逐策略，这还会变得更重要。因此，许多策略本身可以通过影响和控制相关节点的引用计数量，同时依然参照引用计数的大小来制定驱逐策略，在未来的扩展性设计中，这将节省大量开发成本。

设计模式总结，SGLang vs vLLM vs Pico-vLLM

	SGLang	vLLM	Pico-vLLM
索引结构	radix tree	hash table	radix tree
存储粒度	page_size=1	block_size=16	block_size=16
索引粒度	token	block	token
匹配精度	token 级	block 边界	token 级，返回前对齐
引用计数层数	1（lock_ref）	1（block ref）	2（lock_ref + logical_ref）
inc/dec 方式	node 指针 + parent chain	hash key 直接索引	node 指针 + parent chain
分裂处理	天然对齐（page=1）	不分裂	需要处理非对齐切分
驱逐单位	叶子节点	单个 block	叶子节点
驱逐策略	LRU	LRU	LRU + lazy deletion

SGLang 的优势：page_size=1 意味着树和存储粒度完全一致，不存在对齐问题。split 后每条边的 token 都有独立的 page，不会出现"半满 block"的归属问题。代价是 page table 更大（每个 token 一个 entry），管理开销更高。

vLLM 的优势：hash table 没有分裂合并，O(1) 查找，实现简单。block 粒度的管理开销最低。代价是匹配精度受限于 block 边界——两个 prompt 在 block 中间分叉时，整个 block 都 miss，浪费已计算的 KV Cache。

Pico-vLLM 的定位：在 vLLM 的 block-level 存储架构上获得 SGLang 的 token-level 匹配精度。代价是需要一个显式的桥接层处理粒度转换和双层引用计数。可以预计，这应该不是性能上最高效的方案，因为增加抽象层级和转换必然引入开销。不过，个人认为，这样引入的内存布局的兼容性可以在其他地方的可扩展性层面发挥重大作用。以我个人的经验而言，至少有一个地方非常需求这个层面的自由度调优兼容性：PD分离和多卡TP通信。在这两者的通信当中，NCCL需要拷贝后发送，而NIXL的零拷贝则需要按块为粒度进行发送，后者在block size过小的情况下会引入严重的launch overhead，从而极大拖累通信性能。因此，block size的自由调整和radix tree风格的前缀匹配的适配是非常重要且有意义的。

Benchmark和性能测试

好不容易把这么复杂的系统设计完之后，不做benchmark，看看Prefix Caching技术到底有何优劣，就太可惜了。设计两个不同长度的Prompt评测情景，对其进行性能分析。

短 Prompt（144 token）

3 种 system prompt（assistant / coder / analyst），每种 3 轮请求。第 1 轮 cold start，第 2-3 轮 warm。

hit_rate = 62.8%

短prompt下GPU prefill只需几毫秒，CPU overhead（Python侧tensor创建、scheduler调度、slot_mapping计算）占主导。prefix cache节省的GPU时间被CPU开销掩盖，加速效果有限。

实际上，对于小模型、非常短的prompt来说，由于前缀匹配引入了额外的CPU overhead的开销，用时不仅不会减少，反而会明显增加。在实测过程中，对于这个量级的prompt，大约会拖慢30~~40%的执行速度。换言之，Prefix Caching本身所引入的额外CPU overhead大约是~~5ms量级左右。如果Prefix Caching能够节省的内存量超过这个量级，那么Prefix Caching的策略就是值得的。

如果反向思考的话，这本身可能暗示了另一种策略的产生：如果对话不够长、模型足够小，那么也许直接通过某种分离策略来bypass正常的Prefix流程，按照之前的naive paged cache方法进入控制流，反而是更好的做法。

长 Prompt（2083 token shared + ~35 token variable）

固定一个 2083-token 的超长 system prompt（技术参考文档），8 个不同的 user query（各 30-50 token）。共享前缀占比约 98.5%。

共享前缀: 2083 tokens
可变部分: 28-48 tokens

 OFF(ms) ON(ms) speedup
Cold (req0): 48.92 41.59 1.18x
Warm (avg 7): 41.17 16.06 2.56x
Warm (best): 41.71 12.08 3.45x

2000+ token 的 prefill 让 GPU 计算成为瓶颈（约 30-40ms），CPU overhead 占比降到 20% 以下。prefix cache 省掉了 98.5% 的 prefill 计算，warm 请求的 TTFT 降到 cold 的 30-40%。

理论上界分析：prefix cache 的最大加速比是 total_tokens / new_tokens ≈ 2083 / 35 ≈ 60x。实际上的加速比只达到3.45x。这实际上是因为CPU侧的overhead实际上占据了相当大的比例，几乎有~5ms之多，这部分是无法消除的，其开销完全不会随prefix cache命中率的提高而减少（tensor创建、block table构造等仍需要执行）。此外，CUDA Graph 有固定的launch overhead，Python GIL下的 scheduler调度会产生额外开销，adopt_blocks 需要更新 GPU 上的 block table同样会引入额外开销，这些都会在我们评测使用的小模型上产生不可逾越的加速比gap上限。阿姆达定律告诉我们，最终的加速比不可能超过这些部分占的时间比例的倒数。考虑到这一点，Prefix Caching已经在职责范围内很好的完成了它的任务了。

总结和心得

到这里，所有Pico-vLLM计划内的所有主要计划内feature就已经都完整实现了，而且对于单卡的性能部分、PD分离部分的延迟抖动和异构并行度的正确性验证、Prefix Caching部分都有了完整的benchmark和Profiling。在工作过程中，也产生了若干个明显可以后续优化的点，包括了1、TP模式中的通信异步化和层间架构的重设计以允许通算重叠，2、PD分离的时候传递的后端替换为NIXL以评测性能变化，3、Scheduler的Prefill的Chunk策略和更具体调度策略的设计，4、Prefix共享block的写时复制（COW）语义的实现，5、GPU to CPU的换入换出的offload驱逐策略的实现，6、驱逐策略和Radix Tree结构的分离解耦，以及其他CPU侧代码的整体性能优化，都是未来可以不断扩充、实现和改善的地方。这些改善将在后续不断的进行，但集中的开发工程本身即到此为止。总体来说，这一个月的集中开发所取得的成果是令人感到满意的。

然后是心得。其实框架做到最后已经有sys的意思了。分层抽象，协同设计，性能和兼容性的取舍，每一个“传统”的sys设计思想，在推理框架里都有明确的体现。例如，怎么在开始动笔之前设计一个抽象层级，和它的精确的语义？这个很重要，语义设计不清就会权责不明，权责不明就产生耦合，产生耦合后续维护性就变差。有时候，这个问题比性能本身还重要。宁可性能损失3%，换取可维护性，大多数时候可能也是值得的。在这次的例子中，对于每一层的设计来说，应该具有哪些函数、其语义是什么，应该管理的东西有哪些、引用计数的来源和去向，预测的断言和边界条件情况，如果能够先尽可能的设计好再动手开发，可能反而会比现在更省心思、更省时间，也更不容易产生按下一个bug，又连带造成另一个bug，最后在修修补补中逐渐缠绕的一团乱麻的情况。在完善的开发流程中，最好都应该在纸面上先仔细想好、确保边界条件思考的尽可能全面再动笔，不仅可以省下后续大量的debug的无用时间，更可以提升最终成品的质量。本质上，在软件工程里，这种思想已经重复了无数次，但说到底绝知此事要躬行。此处教训相当宝贵，是应当牢牢记住的。

Pico-vLLM 开发日志 #13 PD分离（续）同步异步模式

Mon, 13 Apr 2026 22:31:38 +0800

这次迭代完成的是两个feature：异步模式的PD分离实现，以及和异构并行度下的PD分离-TP耦合的设计兼容。这部分做完之后，就只剩下基于Radix-Tree的Prefix Caching特性需要实现了，随后就进入整理、代码优化和交互界面设计阶段。这篇博客主要记录一下同步-异步模式两个的思路和实现，包括生产者-消费者模型；以及基于异构的PD分离-TP并行设计的具体实现。

什么是同步和异步模式

所谓同步模式，指的是"存在一个可序列化的顺序"，即整个请求的生命周期可以被描述为一个严格的线性序列：Prefill → KV Cache传输 → Decode。换言之，每一步不会在完成之前开始下一步，控制流是单一的。

而所谓异步模式，指的是"不存在这样一个全局的串行顺序"，各阶段之间通过显式的信号机制进行协调，而非通过阻塞等待来保证先后关系。Prefill端完成KV Cache的生产后，不需要等待Decode端确认接收就可以立刻开始处理下一个请求；传输过程可以和两端的计算并发进行；Decode端在KV Cache就绪后被通知并开始工作。控制流不再是单一的，而是分裂为多条独立推进的流，它们之间只在必要的数据依赖点上同步。

这两者的区别在于，当每一步不存在严格的对上一步的依赖，且对硬件资源的依赖不相互重叠的时候，即可通过这种方式进行某种意义上的重叠，从而提高性能。否则，两者的区别就不是很大，异步可能因为信号的轮询反而有性能的额外开销。

同步模式

同步模式的设计很简单。它的流程是：请求到达Prefill端，完整执行所有层的Prefill，生成全部层的KV Cache，然后一次性将KV Cache传输到Decode端，Decode端接收完毕后将请求加入Decode调度队列，开始逐token生成。它当然很容易保证和调试正确性，但它毫无疑问对于生产级来说是不可取的。

异步模式

异步模式的设计复杂度则比预想中高很多。它的最主要障碍反而不在于具体的编码实现，而在于设计模式的收发约定：未来还有多少个请求一定会要发送（发送方和接收方不再同步执行，那接收方怎么知道"该准备接收了"，怎么知道“不用再接收了”），以及后端的特性支持情况：保序性、是否支持tag以区分信道。不幸的是，对于这次实现的torch.distribute使用的nccl后端来说，它保序但不支持tag区分。而对于前者，实际上这没有可以简单解决的方法。

具体来说，NCCL的P2P是双边操作：isend必须和对端的irecv配对才能真正开始传输，单方面调用isend并不会把数据推出去。这意味着Decode端必须提前挂好irecv等着，但问题是它不知道Prefill端什么时候会发、甚至不知道还有没有下一个请求要发。

一个请求的KV Cache传输至少需要三次P2P通信：size、meta、data。在不知道后续是否还有请求的情况下，Decode端无法为下一个请求提前挂出irecv。这直接导致了一个硬约束：KV传输在请求级别是串行的，同一时刻只能有一个请求的传输流程在进行。 上一个请求的三次握手没有走完，下一个请求的传输就无法开始。

当然，更复杂的设计其实是存在的：比如在消息体中携带后续请求的元信息，或者建立一条独立的控制信道来提前广播传输计划，从而让Decode端可以预先挂出多个irecv实现传输的流水线化。但这已经不是推理引擎层面的工作了，而是进入了高性能通信协议栈的设计范畴，超出了本项目的边界。

NIXL的必要性

前面说了这么多，其实，目前的NCCL的局限已经体现的很明显了。这就自然引出了NIXL的需求：如果我们不需要双边的P2P语义，这个问题就根本不存在了。让ai老师介绍一下它的特长所在：

**NIXL（NVIDIA Inference Xfer Library）**是NVIDIA在GTC 2025上开源的、专门为分布式推理场景下的数据搬运设计的点对点传输库。它和NCCL的本质区别在于通信语义：NCCL的P2P是双边的send/recv模型，发送方和接收方必须配对才能完成传输；而NIXL采用的是单边的read/write模型——一个节点通过带外网络（如TCP或ETCD）导出自己的KV Cache元数据，其他节点拿到元数据后可以直接读写远端的GPU显存，远端甚至不需要在那一刻做任何配合操作。这意味着前面讨论的双边配对约束在NIXL的模型下天然不存在：Decode端可以在任意时刻、对任意数量的请求发起KV Cache读取，不需要和Prefill端一一握手。

在底层传输机制上，NIXL支持RDMA（InfiniBand和RoCE）、GPU-Direct RDMA、GPUDirect Storage、NVMe-oF、POSIX socket、以及S3兼容的对象存储等多种后端，并且会根据硬件配置自动选择最优的传输路径。一个关键的性能特性是，NIXL的传输不消耗GPU的SM资源——数据通过GPU-Direct RDMA直接在网卡和显存之间搬运，不经过GPU的计算单元，因此不会和模型的forward计算争抢算力。在实际benchmark中，NIXL在256KB到1MB的典型KV Cache传输大小上比NCCL快30%到50%。

NIXL的核心实现是C++，通过pybind11提供Python绑定，同时也提供Rust绑定。它目前已经被集成到NVIDIA Dynamo、TensorRT-LLM、vLLM、SGLang、LMCache等主流推理框架中，是NVIDIA推理基础设施栈中负责"搬数据"这一层的标准组件。

在后续实现中，也有计划在实现所有核心feature之后，将其变为PD分离的数据搬运的新后端，并且进行benchmark评测。

生产者-消费者模型

生产者-消费者模型是并发编程中最经典的协作模式之一。它描述的是这样一种场景：一方负责生产数据，另一方负责消费数据，两者通过一个共享的缓冲区解耦，各自以独立的节奏运行。生产者不需要等消费者处理完上一份数据才能继续生产，消费者也不需要等生产者准备好下一份数据才能开始处理。只要缓冲区里有东西，消费者就取；只要缓冲区没满，生产者就放。两者之间唯一的同步点是缓冲区的状态：空的时候消费者等待，满的时候生产者等待。

这个模型对通信机制的要求也很明确：第一，需要一个共享的、线程安全的缓冲区；第二，需要一种通知机制让双方知道缓冲区的状态变化（通常是信号量、条件变量或事件）；第三，生产者写入缓冲区的操作和消费者从缓冲区读取的操作不应该要求对方的主动配合。换句话说，理想的通信模式是单边的，生产者侧和消费者侧不需要握手，也不需要提前约定通信的固定总量。

PD分离就是典型的生产者-消费者场景。Prefill端生产KV Cache，Decode端消费KV Cache，两者的处理速度天然不同（Prefill是compute-bound的大块计算，Decode是memory-bound的逐token生成），需要解耦才能各自高效运行。但从实现上看，能否真正实现生产者-消费者模型，取决于底层通信机制是否满足上述要求，特别是第三点：单边性。如果传输机制要求双边配对（如NCCL的isend/irecv），那生产者每次放数据都需要消费者同时伸手来接，这就把生产者-消费者退化成了一个双方必须同步握手的模型，缓冲区的解耦作用被大幅削弱。而如果传输机制支持单边读写（如RDMA），生产者只需要把数据写入一块预先注册的远端内存，消费者在任意时刻来读就行，这是生产者-消费者模型真正能发挥作用的前提。

基于异构的PD分离-TP并行设计

这部分的思想同样也很简单，就是让P侧占用的卡数量（并行度）可以和D侧不同，两者解耦。这么做的好处显而易见如果P和D的计算特性差异很大，那么，给两者分配一模一样的GPU数量很可能是不明智的。此外，它也会影响在特殊集群或者有掉卡情况下集群的弹性扩缩容。而如果P侧和D侧的并行度可以不同，那么整个性能调整就可以很自由，可以通过精细的调优最大化集群硬件的利用率。

KV Cache的重分片

KV Cache跨TP重分片的实现是核心的开发需求。问题的根源在于：TP并行下，每张卡只持有KV heads的一个子集。以Qwen2.5-1.5B为例，模型有2个KV heads，如果Prefill端TP=2，每张卡各持有1个KV head的Cache；而如果Decode端TP=1，单卡需要完整的2个KV heads。这意味着KV Cache的传输不再是简单的整块搬运，而是必须在传输过程中完成一次重新拼合或拆分——从Prefill端的多张卡上各取一部分，聚合成Decode端所需要的完整形状，或者反过来将完整的Cache切分后分发到Decode端的多张卡上。

设模型共有 $N_{kv}$ 个KV heads。在TP并行度为 $tp$ 的情况下，rank $r$ 持有的KV heads范围是：

$$ \text{heads}(r, tp) = \left[\frac{r \cdot N_{kv}}{tp},\;\frac{(r+1) \cdot N_{kv}}{tp}\right) $$

因此，Prefill端rank $r_p$（并行度 $tp_p$）持有的heads范围是：

$$ H_p(r_p) = \left[\frac{r_p \cdot N_{kv}}{tp_p},\;\frac{(r_p+1) \cdot N_{kv}}{tp_p}\right) $$

Decode端rank $r_d$（并行度 $tp_d$）需要的heads范围是：

$$ H_d(r_d) = \left[\frac{r_d \cdot N_{kv}}{tp_d},\;\frac{(r_d+1) \cdot N_{kv}}{tp_d}\right) $$

重分片的本质就是计算 $H_p(r_p)$ 和 $H_d(r_d)$ 之间的交集。对于Decode端的某个rank $r_d$，它需要从所有满足 $H_p(r_p) \cap H_d(r_d) \neq \emptyset$ 的Prefill端rank收集数据。具体来说，$r_d$ 需要接收数据的Prefill端rank范围是：

$$ r_p \in \left[\left\lfloor\frac{r_d \cdot tp_p}{tp_d}\right\rfloor,\;\left\lfloor\frac{(r_d+1) \cdot tp_p}{tp_d} - 1\right\rfloor\right] $$

而从每个这样的 $r_p$ 处，$r_d$ 需要取出的head局部索引范围（相对于 $r_p$ 自己持有的heads的起始位置）是：

$$ \text{local\_start} = \max\left(0,\;\frac{r_d \cdot N_{kv}}{tp_d} - \frac{r_p \cdot N_{kv}}{tp_p}\right) $$

当 $tp_p = tp_d$ 时，所有公式退化为 $r_p = r_d$ 的一对一映射，不需要任何重分片——这也是同构PD分离的特殊情况。

写成公式就是这么几行，但翻译到代码里就会发现：每个传输操作的source rank、目标rank、在KV Cache tensor上的slice范围、以及block_table中的物理块重映射，全都要从这些公式算出来，而且必须在每种 $tp_p$、$tp_d$ 的组合下都正确。这还只是考虑了p、d互相的差异是整数倍的情况。如果不是整数倍、或者并行度大于KVhead的数量，其复杂度将会加倍。非整数倍意味着某个Decode端rank需要的heads会横跨某个Prefill端rank持有的heads区间的中间位置，产生不对齐的切片；而并行度大于KV head数量则需要引入GQA的head重复映射，同一个物理KV head的Cache要被复制到多张卡上。在pico-vllm的实现中，默认p和d的分配一定是整数倍的。

此外有意思的是，vLLM的异构并行度PD分离支持长期处于实验状态，相关的bug和兼容性问题从2024年底的路线图规划到2025年下半年仍在持续修复中。也许从中可以一窥跨并行度KV重映射在边缘情况下的实现难度有多令人头疼。

Pico-vLLM 开发日志 #12 PD分离

Sun, 12 Apr 2026 21:10:14 +0800

PD分离可以说是最近几年最为火热的llm推理引擎架构创新了。从Meta到Openai到Anthropic，内部几乎都肯定有完整的PD分离框架，其中Meta甚至因为这个发过相关的技术论文。它的核心卖点在于推理过程中的一个计算特性区分：Prefill是一个计算强度和prompt length相关的操作，是计算密集的，同时是单步重负载的；而Decode是一个计算强度几乎为固定值的操作，同时也是单步相对轻负载的。在前面的设计我们知道，基于Continuous Batching技术+Chunked Prefill会引入Prefill和Decode的混编，而这种混编会引起TTFT的抖动，以及其他方面的性能影响。如何解决它？答案意外的简单，把它们物理上分开不就行了？

但有意思的是，vllm直到最近还在把PD分离作为一个有experimental标签的“实验性”特征，25年8月才修好P侧和D侧的TP异构导致的不兼容问题。对于一个如此热门的开源项目来说，似乎进展相当慢了。所以在做之前的当时就已经在怀疑，是否可能有坑？做到目前来看的结果，确实是相对来说复杂度较高的feature，原因主要不是技术原理，而是工程实现的复杂度，尤其是和其他并行模式的耦合的复杂度。下面主要记录一下这部分产生的心得。

为什么需要PD分离

这个问题不难理解，在前面也讲到了，本质是计算强度的不同，导致的最优化的ridge point完全不同，进而导致硬件资源的非充分利用。不过在阅读论文之后，我发现其实还有一个原因，就是用户端的体验问题。下面将会逐一讲解。

TTFT和ITL的权衡

TTFT（Time To First Token）衡量的是从用户输入prompt开始，到模型开始生成第一个Decoding Token的时间间隔。这本质上是一个人机工程学指标：毕竟硬件的效率就摆在那里，TTFT是一个“权衡”，而不是算力上限的最精确衡量。但是它对于用户来说是非常重要的：我们在TTFT时间之后才可以真的开始阅读LLM生成的内容、进入非闲置的工作状态，在此之前是纯粹无效等待。它也很可能影响“系统是不是卡了”的直觉判断，一个TTFT高或者不稳定的LLM框架即使后续生成很快，也会无法高效维持用户的注意力。

ITL（Inter-Token Latency）是生成过程中相邻两个token之间的间隔。它决定的是"阅读体验是否流畅"。它同样会影响到阅读体验：如果大量Prefill突然加入，ITL就一定会升高（这是不可避免的），从而让阅读被打断。更重要的是人在阅读情况下对ITL比起TTFT更敏感，因为人没有被打断的心理预期。

这两者之间构成trade-off。如果为了保护ITL，限制每步最多只做一个Prefill（或者限制Prefill的token数），那TTFT就会变长，因为新请求要排队等。如果为了降低TTFT，激进地插入Prefill，那正在Decode的请求的ITL就会被打断。这是一个调度层面无法根治的矛盾，其问题在于“不稳定性”和“不可预测性”。只要Prefill和Decode在同一张卡上，这个跷跷板就存在。

在Continuous Batching框架里，Prefill和Decode共享同一个GPU。当一个新请求进来做Prefill的时候，它的计算量远大于同batch里正在Decode的那些请求（一个Prefill可能处理几百上千个token，而每个Decode只处理1个token）。这意味着这一步的总耗时被Prefill主导，所有正在Decode的请求都必须等这一步跑完才能拿到自己的下一个token。结果就是：正在生成中的请求的ITL被新请求的Prefill拖长了。确保它的稳定性在非分离情况下需要非常非常复杂的调度策略，几乎接近于时序预测而非确定性的调度策略，这就自然催生PD分离需求。

硬件资源的充分利用

更深层次的原因是Prefill和Decode的计算特征差异大到它们的最优硬件配置完全不同，放在一起跑意味着任何时刻都有资源在浪费。正如前面说的，它们具有不同的总体算术强度。我们都知道，硬件的最优化利用是当程序处于其ridge point，或者接近其ridge point的时候。但一个硬件的ridge point显然只有一个，而Prefill-Decode两个不同需求有两个在Roofline模型里差异极大的算术强度点（具体可以参加前面的博客）。当Prefill在跑的时候，HBM带宽没有被充分利用（compute-bound，Tensor Core和CUDA Core是瓶颈）。当Decode在跑的时候，GPU的算力没有被充分利用（bandwidth-bound，HBM是瓶颈）。这意味着无论在哪个阶段，都有一部分的硬件能力在闲置。

PD分离带来一个潜在机会。如果我们进行恰当的分离，是否可以部署异构硬件，从而以相同的预算成本做相同的事情？Prefill用算力更强的卡，Decode用显存带宽更高或者成本更低的卡。实际上，最近几年Nvidia和华为也都在基于Prefill和Decode阶段的不同计算特性，设计具有不同ridge point的GPGPU（Rubin CPX）或者NPU硬件变体。它们基本上也就是为了匹配PD分离里分离不同算术强度的操作而设计的，避免资源闲置和同成本下的利用效率最大化。

即使是在软件层面，也已经有了（实际上更完善的）软件PD分离支持，这部分可以参考NVIDIA Dynamo（官方的PD分离调度框架）和NIXL（Inference Xfer Library），在此不过度展开了。

可用的工具箱

对于具体的PD分离实现，其实有很多种不同选择，以下几种均有实现的潜力。为了单人开发的复杂性和可维护性控制，笔者选择的是直接使用torch.distribute。然而，此处也列出工业界和其他可用的选项，供大家参考。这部分很大程度上是通过ai提供思路，然后再进行调研的，再次感谢ai。

使用torch.distribute（NCCL）

这是最容易上手的方案，也是之前Pico-vLLM在TP的开发阶段已经用过的工具。torch.distributed封装了NCCL的集合通信原语，提供了send()/recv()的点对点通信接口。

优点：开发成本低。NCCL对NVLink、PCIe、InfiniBand的底层拓扑发现和路径选择都是自动的，不需要开发者手动管理。

缺点：NCCL的设计初衷是集合通信，而不是点对点的大块数据传输。另一个问题是NCCL要求通信的所有进程在初始化时就加入同一个通信组（process group），这意味着Prefill Worker和Decode Worker必须在启动时就互相知道对方的存在。这在静态部署下没问题，但如果想要做动态的弹性扩缩容（比如根据负载动态增减Prefill/Decode Worker的数量），NCCL的这种静态组模型就会变得非常不灵活。此外，如果想在这个方案下做，需要先把传输内容gather到一起，进行传输，然后再在远端scatter。这意味着两次额外拷贝。

使用NIXL

NIXL是NVIDIA在GTC 2025上开源的、专门为推理场景设计的点对点数据传输库。它是NVIDIA Dynamo推理框架的核心数据搬运层。和NCCL不同，NIXL从设计之初就是为了解决推理场景下的KV cache传输这个具体问题的。

NIXL的核心抽象是Agent和可插拔的传输后端。每个Worker（无论是Prefill还是Decode）运行一个NIXL Agent。Agent向NIXL注册自己的内存区域（GPU HBM、CPU DRAM、NVMe SSD等），然后通过异步的transfer()接口发起传输。NIXL会自动选择最优的传输后端。

优点：专为推理KV cache传输设计；支持异步、非阻塞传输；支持非连续内存scatter/gather；多后端自动选择；与vLLM、SGLang、TensorRT-LLM等主流框架已有集成。

缺点：依赖较重（需要UCX、RDMA驱动等），在没有RDMA的环境下（比如我们的autodl PCIe环境）退化为TCP。NIXL目前仍在快速迭代中，API稳定性和文档完善度还有提升空间。安装也不是特别顺畅。此外更重要的是，它是C原生的，将其整合进入Python需要额外的大量努力。

使用CUDA IPC

CUDA IPC是CUDA原生提供的跨进程GPU内存共享机制。核心API是cudaIpcGetMemHandle()和cudaIpcOpenMemHandle()：前者从一块已分配的GPU内存中导出一个可序列化的handle，后者在另一个进程中用这个handle获取一个指向同一块物理显存的指针。

优点：同机部署下理论上是延迟最低的方案。因为它是零拷贝的，传输延迟为零。实现也相对直接，不需要额外的库依赖。

缺点：只能在同一台机器上使用。这是因为它需要同一块物理显存。如果要通过网络，自然就没有所谓的共享虚拟内存空间了。

设计模式

基于实例角色role的兼容式设计

这部分的想法很简单：在不额外增加函数的情况下实现不同实例角色的指派，而在实现上保持兼容。

实例具有三种可能的角色：“p”、“d”和“pd”。

传输层抽象分离

一种简单的想法是把传输实现在Engine里。它是最简单的，但思考之后我并没有采纳这样的方案。原因是传输的具体方案实在是太多了：就像刚刚说的那样，不同的后端实现、同步异步、和Scheduler策略的耦合、P侧和D侧的收发约定将会带来不可维护的复杂性。因此不如单独抽象一个基类，然后作为接口来开发。这是软件工程的想法，但值得一提。目前实现的是同步模式，但是，传输层的分离将允许未来进行异步和更加复杂的可扩展性。

注意事项和Bug

值得注意的Bug汇总

这里总结了几个遇到的bug，详细记录一下，以免后面忘记再踩重复的坑，同时也给后来的读者一个提示，避免有相同的错误。

1、Python的可变默认值问题

在开发过程中，遇到类似下面的bug。当提交若干个（超过1个prompt的时候），出现了具有奇怪但是固定模式的输出乱码。如下图：

 [Request 1] The capital of France is Paris Paris 11200
AAExample 112

 [Request 2] 1 + 1 = Paris Paris 11200
AAExample 112

输入的prompt分别是：

1
2
3
4
5


prompts = [
 "The capital of France is",
 "The capital of France is",
 "1 + 1 =",
]

其中第一个request的结果被彻底吞掉，而。这个我一开始以为是KV cache相关的部分又出了问题，因为太像了。经过长达三个小时的详细排查，最终发现了一个bug的特性：它重复的Paris这个首个Decode token的数量，和input的prompt完全相关。当输入的prompt变成：

1
2
3
4
5
6


prompts = [
 "The capital of France is",
 "The capital of France is",
 "The capital of France is",
 "1 + 1 =",
]

的时候，错误的输出结果变成了：

[Request 3] 1 + 1 = Paris Paris Paris 1112000
AAAExample

又多了一个Paris。这给了我缩小范围的想法：它很可能是很多个prompt的Decode结果，同时挂在了所有的prompt的generated_ids序列上。那么，它的工作机制是什么呢？先看原始request类的代码：

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


class RequestStatus(Enum):
 WAITING = "waiting" # 在 waiting 队列，还没做 prefill
 PREFILL = "prefill" # 本步正在做 prefill
 DECODING = "decoding" # 已经 prefill 完，正在 decode
 FINISHED = "finished" # 已完成

''' 请求对象，包含请求的所有信息和状态
- request_id: 请求 ID，唯一标识一个请求
- input_ids: 输入的 token ids，shape (1, init_seq_len)，包含整个 prompt
- generated_ids: 已经生成的 token ids，shape (1, )，初始为 空
- max_new_tokens: 最多生成多少个 token
- temperature: 采样温度，传递给 sampler
- top_p: top-p 截断，传递给 sampler
- kv_cache: 每个请求独享一个 KV cache 实例，存储生成过程中的 KV 状态
'''
class Request:
 request_id: int
 input_ids: List[int] # (1, init_seq_len)，包含整个 prompt
 generated_ids: List[int] # (1, )，包含已经生成的 token ids，初始为 空
 max_new_tokens: int
 temperature: float
 top_p: float
 kv_cache: PagedKVCache # 每个请求独享一个 KV cache 实例
 request_status: RequestStatus
 has_eos_token: bool # 是否已经生成 eos_token，scheduler 不直接接触 tokenizer 和 eos_token_id，这个由 engine 在 decode_step 后更新

 def __init__(self, request_id: int, input_ids: List[int], max_new_tokens: int, temperature: float, top_p: float, kv_cache: PagedKVCache, generated_ids: List[int]=[]):
 self.request_id = request_id
 self.input_ids = input_ids
 self.generated_ids = generated_ids
 self.request_status = RequestStatus.WAITING
 self.max_new_tokens = max_new_tokens
 self.temperature = temperature
 self.top_p = top_p
 self.kv_cache = kv_cache
 self.has_finished_notification = False # engine改变这个状态，scheduler根据这个状态改变 request_status和移出队列

 def is_max_len_finished(self) -> bool:
 return len(self.generated_ids) >= self.max_new_tokens

 @property
 def prompt_len(self) -> int:
 return len(self.input_ids)
 
 @property 
 def total_len(self) -> int:
 return len(self.input_ids) + len(self.generated_ids)

观察这一句：

1

def __init__(self, generated_ids: List[int] = []):

它涉及到一个Python语法：Python的函数默认值在函数定义时求值一次，不是每次调用时。[] 在 Python 解释器加载这个类定义的时候就创建了一个 list 对象，存在 Request.__init__.__defaults__ 里。之后每次调用 Request(...) 不传 generated_ids 时，用的都是同一个 list 对象。所以 r0.generated_ids.append(42) 之后，r1.generated_ids 也变成了 [42]，因为它们指向内存中同一个 list。这是一个非常隐蔽的bug，因为大部分测试原本是在B=1情况下验证，此时只有一个Request实例，共享不共享都无所谓。B>1时多个 Request 往同一个list里append，所有请求的 generated_ids[-1] 返回同一个值，从外部看起来，就是decode出了相同的token。

在修复了Request类中相关的错误代码之后，这个问题就成功解决了，框架开始输出正确的request结果。如图：

 [Request 0] The capital of France is Paris. The capital of the United States is Washington, D.C. The capital of Canada is Ottawa
 [Request 1] The capital of France is Paris. The capital of the United States is Washington, D.C. The capital of Canada is Ottawa
 [Request 2] 1 + 1 = 2
Example 2:
Input: nums = [1,2,3,4,

2、锁死问题

这个严格来说不是实现的bug问题，而是设计模式的问题。当我实现到一半，在写测试脚本的时候，我意识到一个问题：对于一个有并发的项目，接受方怎么确认不用再进行通信了？这个问题乍一看可以通过预先约定次数来解决，但实际上比这个复杂一些，因为很多情况下你根本没有办法提前预知次数，当然也就没法约定它。又考虑了，如果传递一个空请求，那么就自动结束，不再接收信息。这个逻辑在同步下有效，但是在异步下无效：你如何确认在Prefill端没有请求之后，就不会再出现请求了？也许只是短暂出现了波谷，而不是真的停止了。但如果这样不退出的话，那么实际上你没有办法在Engine内部设置一个正常的退出条件，只能通过狼狈的ctrl+c来强行终止。思来想去，最后的结论是：无法在Engine内部做出这个判断，它应该是用户层控制的部分。

为此进行了重构，现在，Engine在用完之后，用户需要显式调用一次“no more requests”函数，以此显式通知Prefill侧的Engine不再接受新request的submit，并且在清空流水线之后通知Decode侧，让它在清空流水线之后同样退出。这样，就比较好的解决了锁死问题。

Pico-vLLM 开发日志 #11 张量并行（续）数据对比分析

Sat, 11 Apr 2026 20:32:26 +0800

正如之前所说，在b200上用单卡和两卡重新测试了相同的1.5B模型的结果，进而得到了一些有意思的小结论。数据删去繁杂的细节之后（感谢ai），整理在下表中。

性能数据汇总

wall-clock端到端时间（含CPU开销）

硬件 tp 延迟 tok/s 加速比
───────────────────────────────────────────────────────────────────────
5090 PCIe (Py3.12+Xeon 8470Q) 1 2.74ms 365 —
5090 PCIe (Py3.12+Xeon 8470Q) 2 2.13ms 470 1.29x

B200 NVLink (Py3.10+Xeon 6960P) 1 1.64ms 610 —
B200 NVLink (Py3.10+Xeon 6960P) 2 2.08ms 480 0.79x

可以看到，前者可以达到一定程度的加速，而后者不仅没有加速，反而被拖慢了。这可能是什么原因？第一反应肯定是因为B200本身算力太高，分摊GEMM计算量的好处相当有限，而通信带来的提升不足以被PCIe到NVLink的配置升级所补偿。算力的升级幅度远大于通信的升级幅度，因此前者可以加速，而后者就不行了。

这个解释本身就是可以接受的。不过，还有一个疑惑：NVLink的性能比PCIe按理来说高很多，B200的性能，无论是算力还是带宽，按理来说也比5090高很多，但为什么在tp=2的情况下，两者的延迟如此接近，几乎完全相等（只有不到3%）？接下来的分析揭示了对小模型（以及不够大的大模型）来说另一个很隐蔽，但是可能同样也很重要的开销：CPU调度开销。

纯GPU replay时间（无CPU开销）

为了解答这个问题，设计了纯基于CUDA Graph的连续replay测试。这部分测试是为了揭示在没有CPU控制流开销的情况下GPU本身的性能如何，以反过来倒推CPU开销的大小如何对性能产生影响（这个的测试没有那么直截了当）。

见下表。可以看到一个非常有趣的细节：在5090+PCIe的配置组合上，端到端的开销和纯GPU开销相对来说较小：CPU侧的开销要么很小，要么可以被部分的overlap住。而对于B200+NVLink来说，其平均的CPU侧开销几乎是5090+PCIe的两倍以上。这主要有这样几个原因：

1、CPU主频不同。CPU侧的代码并没有那么高的并行度，堆叠核数的意义小很多。而Xeon 8470Q的睿频是约3.8GHz，与此同时Xeon 6960P的单核睿频是约3.2GHz。这两者的性能差异造成了很大的CPU侧时延不同，因此最终在并行度增加的时候抹平了GPU侧的性能优势。

2、Python版本的不同。更新的版本具有更完善的解释器优化，因此能够更好的降低CPU侧阻塞时间。

因此，CPU在现代大模型推理里真的不重要吗？其实不完全是。在模型较小、或者并行度够高以至于单步推理时间较为短的时候（以我们的测试为例，小于2.0ms的时候），CPU的overhead并不一定总是能够忽略的。因此，CPU侧代码的优化和恰当实现同样有其必要性。

硬件 tp 纯GPU延迟 GPU加速比
────────────────────────────────────────────────────
5090 PCIe 1 2.62ms —
5090 PCIe 2 2.12ms 1.24x

B200 NVLink 1 1.46ms —
B200 NVLink 2 1.89ms 0.77x

AllReduce 延迟（单次3KB传输）

可以看到，NVLink的单次通信固定延迟大约是PCIe的三分之一，而带宽是20倍，这也是一个有助于建立直感的数字。它也再次验证了计算机学科几乎通用的一个真理：降低延迟永远比提高带宽困难多了。小数据量的AllReduce是纯latency-bound，其对性能造成的损害于取得的收益相比不成比例（至少不是线性的），因此只适合在足够大的模型、足够好的配置下才值得用，尽管它是实现起来最直接也最方便的。

 固定延迟 数据传输 总延迟 56次/步
─────────────────────────────────────────────────────────────────
5090 PCIe ~28μs 0.06μs ~28μs 1.59ms
B200 NVLink ~10μs 0.003μs ~10μs 0.57ms

值得一提的细节

benchmark环境的隐藏变量不一定可以忽略，有时候影响显著。在最初测试的时候，B200 NVLink (Py3.10+Xeon 6960P)的配置在tp=2的时候甚至慢于5090 PCIe (Py3.12+Xeon 8470Q)的tp=2的情况，而这几乎说不通。经过检查，发现是环境变量中，OMP_NUM_THREADS设置不同（B200侧未指定，默认为1），autodl平台进行了专门的预先设置（指定为了50）但B200的环境中没有，这造成了明显的性能改变，通过拖累CPU侧的性能进而拖累了整个端到端的性能。

此外，这也验证了另一个问题：TP不是万能药，甚至有时候有害。对于小模型、网络延迟不够低的机器配置，尤其如此。

Pico-vLLM 开发日志 #10 张量并行 TensorParallel

Thu, 09 Apr 2026 23:19:56 +0800

在之前的单卡阶段结束和进行阶段性的总结之后，经过几天的休息（实际上是作者被导师叫去忙活论文投稿和下一篇论文启动的事情了），我们终于可以进入Pico-vLLM开发的下一个阶段：分布式的开发。而分布式的开发，根据计划，我们将要实现如下几个重要特性：张量并行TP，流水线并行PP和PD分离。

首先做的是张量并行，Tensor Parallel。其实我在一开始规划的时候，曾经想过要不要先做流水线并行。然而经过进一步的思考和分析，我发现了一个令人惊讶的事实：张量并行这个看上去切分粒度更细、通信要求和频率更高的方案，在实现难度上居然是最低的，而pipeline并行这种方案，需要的改动反而更大。下面会详细解释实现所涉及到的改动，以及实现的具体技术手段和注意事项、细节。

张量并行涉及到的改动

张量并行涉及到的改动意外的少，这个少是相较于其他并行方案而言。这其实是因为张量并行并不改变任何一张单个GPU上所需要加载的模型的结构：所有层数没有变，block没有变，变化的只有Q head和KV head的头数。

模型

模型的结构并不需要修改，仅仅将num_heads相关的参数和hidden_dim相关的参数修改成相应的local版本即可，层数等等结构相关的都不需要改变。

KV cache的存储

和先前的模式一模一样，每个GPU对应的程序仍然要单独管理自己的KV cache，有自己的Scheduler。区别在于，每个KV cache存储的slot大小变化了，因为每个GPU对应的head数量变化了（被分给了不同的GPU），其中缩小的倍率和TP的size相同。除此之外，几乎没有任何需要修改的地方，包括Triton Kernel。

引擎Engine

几乎没有需要修改的地方，只需要在初始化的时候给cfg传入相应的额外参数即可。

另一个需要注意的事情是，如果采用的是有随机性的采样，需要在每一张卡的程序实例初始化的时候，都进行相同的随机数指定。否则，不同的卡将可能会在Logits生成之后产生不同的采样结果，而整个sequence的生成Token是每个GPU上各自解码、各自保存的，每次Decode步之间并没有广播同步（为了性能），因此这将引发不同GPU间的序列不一致性，导致严重的问题。考虑到整体的设计模式，最好的指定是在初始化Engine实例的时候作为一个参数传入。因为用户不一定想到做这件事而且其具体影响不是很大，最好将其设置为一个可选的默认参数：比如42。

张量并行的实现手段

意外的简单，整个改造实现完全在pytorch内部进行。nccl相关的操作被torch.distribute完全封装，而且同样可以被CUDA Graph捕获，正常执行。

需要注意的技术细节

需要格外说明的是，CUDA Graph和torch.distribute+nccl backend会产生一些奇妙的化学反应。具体来说，CUDA Graph在capture的过程中，如果使用了nccl的后端作为torch.distribute的调用，那么CUDA Graph会直接固定的占用nccl的资源。这会导致在程序结束后，dist.destroy_process_group()无法正确完成，整个集群hang住，无法结束。为了解决这个问题，需要显式的删除CUDA Graph的对象：即执行del g, static_output类似的语句。这是一个说大不大但是很烦人的bug，花了一会儿才定位出来，一开始还以为是哪里同步错误，导致计算流本身有问题。

张量并行在Qwen2.5-1.5B上的效果

效果其实并没有预想中好，主要的原因是——虽然难以启齿——没钱。在2x5090，PCIe连接的的autodl平台上，总体来说加速比只有1.3x，效果并不是非常理想。通过Profiling排查原因可以看到，在有Profiling的情况下，通信（All-Reduce）占据了~48%的时间，几乎就是GPU time的一半了。这主要是因为PCIe连接的Latency问题：每次传输的内容其实只有几KB，数据量本身并不构成时间来源。但是，nccl+pcie的allreduce本身的overhead启动延迟就非常长，每次启动几乎需要26μs的时间，这在28层*2次调用的前向传播来说占据的时间长度几乎和真正的计算和内存存取一样多，从而大大拉低了潜力。这个问题没有什么好办法解决，是完全硬件上的。要解决倒也简单：NVLink，它的通信延迟比PCIe低很多。

画饼

想办法预定到了8卡B200的集群，真是太好了。现在是在autodl上做的，因为是PCIe互联，本身就不是非常适合TP这个并行模式，所以加速效果也有限，更何况只有2卡的5090。过几天马上就上去试试集群，看看Profiling的结果会不会有什么变化？能够预计，它能够达到的有效加速比应该比PCIe的情况高。

LLM 学习日志 #2 集合通信拓扑、原语和实现

Sun, 05 Apr 2026 16:39:11 +0800

这篇博客主要是重温一下之前学过的集合通信的基本模式，包括拓扑、原语和实现。说是“重温”是因为这部分在之前学习NCCL的时候其实就都已经粗略的学习过了，但之后没有用到过，因此掌握不熟练、印象不深。现在趁着Pico-vLLM的开发机会，把这部分知识彻底的掌握起来。

集合通信的原语

点对点（P2P）

Broadcast

Broadcast意味着一对多，一个GPU的数据广播给所有人。

GPU 0: [A B C D] → GPU 0: [A B C D]
GPU 1: [ ] GPU 1: [A B C D]
GPU 2: [ ] GPU 2: [A B C D]
GPU 3: [ ] GPU 3: [A B C D]

Reduce

Reduce意味着多对一，所有GPU的数据归约到一个GPU。要注意的是它不是拼接而是运算，其最终得到的数据的数据量和原本每个GPU自己的数据量是相同的。Reduce本身并不是一个具体的操作，而是一系列具有特殊性质的规约算子的集合，其满足的性质是为可结合（associative）且通常可交换（commutative）的二元运算。Reduce的常见算子包括了：

Sum（求和）
Prod（乘积）
Max / Min（极值）
Avg（平均）
BAND / BOR / BXOR（按位运算）
MinLoc / MaxLoc（极值+极值所在位置）

LLM训练中99%的Reduce操作都是求和，其他算子出现的极少。NCCL甚至不支持MinLoc / MaxLoc操作。

GPU 0: [A₀] → GPU 0: [A₀+A₁+A₂+A₃]
GPU 1: [A₁] GPU 1: [ ]
GPU 2: [A₂] GPU 2: [ ]
GPU 3: [A₃] GPU 3: [ ]

All-Reduce

All-Reduce意味着将所有GPU的数据归约，结果每个GPU都有一份完整拷贝。从逻辑上可以理解为先完成一个Reduce再完成一个Broadcast——实际上的操作在概念上类似但具体实现并不相同，同样是分解为两个阶段，但具体分解成的是Reduce-Scatter和All-Gather。

为什么不是一个Reduce+一个Broadcast？实际上，在早期的参数服务器模式下，确实就是如此。但它有一个问题：可扩展性。Reduce 阶段所有数据汇聚到一个节点（CPU），Broadcast 阶段再从这一个节点（GPU）发出去。这个节点（GPU）的带宽成为整个操作的瓶颈，其他 GPU 在等待时空闲。Reduce-Scatter+All-Gather模式实际上是通过一个线性的偏移把每一个节点（GPU）变成了一个$1/N$的参数服务器，每个只负责和自己编号相同的第$K$块的数据的Reduce和Broadcast。在集合通信的视角下，这恰好就是Reduce-Scatter+All-Gather的模式。这样讲，可能在概念上是最容易理解的。

GPU 0: [A₀] → GPU 0: [A₀+A₁+A₂+A₃]
GPU 1: [A₁] GPU 1: [A₀+A₁+A₂+A₃]
GPU 2: [A₂] GPU 2: [A₀+A₁+A₂+A₃]
GPU 3: [A₃] GPU 3: [A₀+A₁+A₂+A₃]

Reduce-Scatter

Reduce-Scatter操作意味着归约后，结果被切分，每个 GPU 只拿到一部分。但这种简化的描述方式实际上不完全准确，因为理解它在过程中的副产物，比理解它的开始状态和结束状态更有助于理解它是如何运作的。以最经典的ring为例子讲解，一张图会更容易帮助读者理解到底发生了什么：

每步：每个 GPU 向下一个 GPU 发送一块数据，同时从上一个 GPU 接收一块并求和。
规则：第 1 步发送自己编号对应的块，之后每步发送刚更新过的块。

初始状态:
 GPU 0: [A₀ B₀ C₀ D₀]
 GPU 1: [A₁ B₁ C₁ D₁]
 GPU 2: [A₂ B₂ C₂ D₂]
 GPU 3: [A₃ B₃ C₃ D₃]

第 1 步: GPU 0 ──A₀──→ GPU 1
 GPU 1 ──B₁──→ GPU 2
 GPU 2 ──C₂──→ GPU 3
 GPU 3 ──D₃──→ GPU 0

 GPU 0: [A₀ B₀ C₀ D₀+D₃ ]
 GPU 1: [A₁+A₀ B₁ C₁ D₁ ]
 GPU 2: [A₂ B₂+B₁ C₂ D₂ ]
 GPU 3: [A₃ B₃ C₃+C₂ D₃ ]

第 2 步: GPU 0 ──D₀₃──→ GPU 1 （发送刚更新的 D 块）
 GPU 1 ──A₀₁──→ GPU 2 （发送刚更新的 A 块）
 GPU 2 ──B₁₂──→ GPU 3 （发送刚更新的 B 块）
 GPU 3 ──C₂₃──→ GPU 0 （发送刚更新的 C 块）

 GPU 0: [A₀ B₀ C₀+C₂₃ D₀₃ ]
 GPU 1: [A₀₁ B₁ C₁ D₁+D₀₃ ]
 GPU 2: [A₂+A₀₁ B₁₂ C₂ D₂ ]
 GPU 3: [A₃ B₃+B₁₂ C₂₃ D₃ ]

第 3 步: GPU 0 ──C₀₂₃──→ GPU 1 （发送刚更新的 C 块）
 GPU 1 ──D₀₁₃──→ GPU 2 （发送刚更新的 D 块）
 GPU 2 ──A₀₁₂──→ GPU 3 （发送刚更新的 A 块）
 GPU 3 ──B₁₂₃──→ GPU 0 （发送刚更新的 B 块）

 GPU 0: [A₀ B₀₁₂₃=B_all C₀₂₃ D₀₃ ]
 GPU 1: [A₀₁ B₁ C₀₁₂₃=C_all D₀₁₃ ]
 GPU 2: [A₀₁₂ B₁₂ C₂ D₀₁₂₃=D_all]
 GPU 3: [A₀₁₂₃=A_all B₁₂₃ C₂₃ D₃ ]
规约完成。

因此，它的最终结果就是下图。但需要注意的是，它并不是一个没有副产物的操作。如果是原位操作的话，它实际上会修改每个GPU没有得到最终规约结果的相应数据位置，使得它不再是原本值，而是一个和自身rank决定的偏移量、数据偏移量都相关的部分累加和，也可以和某种意义上的前缀和相类比。

GPU 0: [A₀ B₀ C₀ D₀] → GPU 0: [A₀+A₁+A₂+A₃]
GPU 1: [A₁ B₁ C₁ D₁] GPU 1: [B₀+B₁+B₂+B₃]
GPU 2: [A₂ B₂ C₂ D₂] GPU 2: [C₀+C₁+C₂+C₃]
GPU 3: [A₃ B₃ C₃ D₃] GPU 3: [D₀+D₁+D₂+D₃]

All-Gather

每个GPU有一部分数据，收集后每个GPU都有完整数据。Ring All-Gather中没有任何归约（求和）操作，每个GPU拿到的块直接原样转发给下一个人就行。每步每个GPU把自己最近收到的完整块沿环传递，N-1步之后每个GPU就收集齐了所有块。

GPU 0: [A] → GPU 0: [A B C D]
GPU 1: [B] GPU 1: [A B C D]
GPU 2: [C] GPU 2: [A B C D]
GPU 3: [D] GPU 3: [A B C D]

All-to-All

每个 GPU 向每个其他 GPU 发送不同的数据。由于其$O(n^2)$的通信连接数（有时候还有通信量）和复杂度，其被称为是“集群性能的试金石和考验”。它几乎总是可能引起拥塞。另一方面，大部分运行在集群上的应用都会尽量避免产生频繁的All-to-All通信，因为它很可能严重拉低相较于规整的集合通信模式下的集群性能。对All-to-All的通信优化算法和工程实践思路以处理其产生的拥塞控制为主。此外，就像上一篇的博客中对于MoE的介绍一样，许多All-to-All是动态的而非静态的，从而进一步加剧了这个问题。

GPU 0: [A₀ A₁ A₂ A₃] → GPU 0: [A₀ B₀ C₀ D₀]
GPU 1: [B₀ B₁ B₂ B₃] GPU 1: [A₁ B₁ C₁ D₁]
GPU 2: [C₀ C₁ C₂ C₃] GPU 2: [A₂ B₂ C₂ D₂]
GPU 3: [D₀ D₁ D₂ D₃] GPU 3: [A₃ B₃ C₃ D₃]

All-to-All-v

每个 GPU 向每个其他 GPU 发送不同的数据，但收发的总数据量不一定每个都相同。

集合通信的常见拓扑

Ring（环形）

最经典的集合通信的拓扑模式。

优点是带宽最优——每个 GPU 的发送和接收带宽被充分利用，总通信量和 GPU 数无关（都是 2×(N-1)/N × 数据量）。

缺点是延迟随 GPU 数线性增长——步数是 2(N-1)，GPU 越多步数越多。

Tree（树形）

优点是延迟最优——只需 O(log N) 步。

缺点是带宽利用率不如 Ring——叶节点只参与部分通信，根节点成为瓶颈。

Double Binary Tree

两棵互补的二叉树同时工作，结合了Tree的低延迟和更好的带宽利用率。NCCL在某些场景下会用这种策略。但其复杂且难以理解，而且对集群的节点数量有更严格的要求。

Recursive Halving-Doubling

递归地把GPU分成两半，先在小组内归约，再逐步扩大。兼顾延迟和带宽。

集合通信的常见实现

XCCL：NCCL、HCCL、ACCL...许多初创的CCL通信库实际上都是直接魔改的NCCL。

NCCL 的 kernel 是沿三个维度预编译的模板实例：

算法（Algorithm）：决定数据流动的拓扑模式

Ring: 环形，带宽最优
Tree: 树形（双二叉树），延迟最优
CollnetDirect / CollnetChain: 利用 SHARP 等网内计算

协议（Protocol）：决定数据传输的同步和打包方式

Simple: 大块传输，高带宽，适合大消息
LL: Low Latency，8 字节粒度，用 flag 做同步，适合小消息
LL128: 128 字节粒度，利用 NVLink 的原子操作，延迟和带宽的折中

归约操作 × 数据类型：

Sum/Prod/Max/Min × FP16/BF16/FP32/FP64/INT8/...

这些维度的组合在编译时就被实例化成了大量的 kernel 变体。例如 Ring + Simple + Sum + FP16 是一个 kernel，Tree + LL128 + Max + FP32 是另一个 kernel。这也是为什么 NCCL 的编译时间很长。

Cost Model：运行时选择

NCCL 的 cost model 是默认调优决策的核心。这个模型以时间为指标评估集合操作的开销，用于选择正确的协议和算法。cost model 考虑多种因素，包括数据量、GPU 架构、拓扑结构、网络和算法属性。

MSCCL：微软用的一种DSL，基于NCCL自定义通信算法。

MPI：太经典不讲。

硬件互联的数感建立

这一节列举常见网络和网卡设备的所有规格和速度，以帮助读者，包括写博客的我自己，建立对于网络的数字感知。

一、GPU 节点内互联：NVLink

代际 首发GPU 每条链路带宽(双向) 每GPU链路数 每GPU总带宽(双向)
NVLink 1.0 P100 40 GB/s 4 160 GB/s
NVLink 2.0 V100 50 GB/s 6 300 GB/s
NVLink 3.0 A100 50 GB/s 12 600 GB/s
NVLink 4.0 H100 100 GB/s (注1) 18 900 GB/s
NVLink 5.0 B200 100 GB/s 18 1,800 GB/s
NVLink 6.0 Rubin 200 GB/s 18 3,600 GB/s

注1: H100 的 NVLink 4.0 每条链路实际是 50 GB/s 双向，但 NVIDIA 在产品规格中
 按 sub-link 口径标注为 900 GB/s（18 sub-links × 50 GB/s）。
 B200 的 NVLink 5.0 SerDes 速率翻倍，同样 18 条但每条 100 GB/s，所以总带宽 1.8 TB/s。

二、GPU 节点内互联：NVSwitch

代际 配套GPU 每芯片端口数 每芯片交换带宽 单节点GPU数 节点总NVLink带宽
NVSwitch 1.0 V100 18 900 GB/s 8 (DGX-2) ~2.4 TB/s
NVSwitch 2.0 A100 36 (注2) ~3.2 TB/s 8 (DGX A100) ~4.8 TB/s
NVSwitch 3.0 H100 64 25.6 Tbps 8 (DGX H100) ~3.6 TB/s
NVSwitch 4.0 B200 72 NVLink5 14.4 TB/s 72 (NVL72) ~130 TB/s (机柜级)

注2: DGX A100 有 6 片 NVSwitch 2.0，每个 A100 的 12 条 NVLink 分别连到 6 片 NVSwitch
 （每片 2 条），实现 8 GPU 全连接。

NVSwitch 3.0 (Hopper) 开始引入 SHARP（Scalable Hierarchical Aggregation and Reduction Protocol），支持在交换芯片内部直接做归约计算（如 All-Reduce 的 sum），数据不需要绕回 GPU。

NVSwitch 4.0 (Blackwell) 把 NVLink 域从 8 GPU 扩展到了 72 GPU（整个机柜），是一个质变——以前跨节点才需要 InfiniBand，现在机柜内全部走 NVLink。

三、CPU-GPU 互联：PCIe

代际 发布年份 每通道速率 x16 单向带宽 x16 双向带宽
PCIe 3.0 2010 8 GT/s ~16 GB/s ~32 GB/s
PCIe 4.0 2017 16 GT/s ~32 GB/s ~64 GB/s
PCIe 5.0 2019 32 GT/s ~64 GB/s ~128 GB/s
PCIe 6.0 2022 64 GT/s ~128 GB/s ~256 GB/s
PCIe 7.0 2025(规范) 128 GT/s ~256 GB/s ~512 GB/s

对比：H100 的 NVLink 总带宽 900 GB/s，而 PCIe 5.0 x16 双向才 128 GB/s，相差 7 倍。这就是为什么 TP 必须走 NVLink 而不能走 PCIe。

四、节点间互联：InfiniBand

代际 发布年份 每通道速率 4x 链路带宽 单/双向 典型网卡 典型交换机
SDR 2001 2.5 Gbps 10 Gbps 单向 - -
DDR 2005 5 Gbps 20 Gbps 单向 - -
QDR 2007 10 Gbps 40 Gbps 单向 ConnectX-2/3 -
FDR 2011 14 Gbps 56 Gbps 单向 ConnectX-3 SwitchX
EDR 2014 25 Gbps 100 Gbps 单向 ConnectX-4/5 SwitchIB-2
HDR 2019 50 Gbps 200 Gbps 单向 ConnectX-6 Quantum
NDR 2022 100 Gbps 400 Gbps 单向 ConnectX-7 Quantum-2
XDR 2025 200 Gbps 800 Gbps 单向 ConnectX-8 Quantum-X800
GDR ~2028 400 Gbps 1.6 Tbps 单向 (规划中) (规划中)

换算成更直观的单位（双向）：

EDR: ~25 GB/s 双向
HDR: ~50 GB/s 双向
NDR: ~100 GB/s 双向 ← 当前主流 AI 集群配置
XDR: ~200 GB/s 双向 ← 2025 年开始部署

一台服务器通常配备多块网卡。例如 DGX H100 配 8 张 ConnectX-7 NDR 400G 网卡，每 GPU 一张，节点间总带宽 = 8 × 100 GB/s = 800 GB/s 双向。即便如此，仍然只有 NVLink (900 GB/s) 的约 89%。

五、节点间互联：以太网（RoCE）

速率 双向带宽 延迟（典型） 备注
10 GbE ~2.5 GB/s ~10-50 μs 传统数据中心
25 GbE ~6.25 GB/s ~5-20 μs 
40 GbE ~10 GB/s ~5-20 μs 
100 GbE ~25 GB/s ~2-10 μs RoCEv2 常见配置
200 GbE ~50 GB/s ~2-5 μs 
400 GbE ~100 GB/s ~1-3 μs 正在部署
800 GbE ~200 GB/s ~1-2 μs 2025 年开始

以太网 + RoCE（RDMA over Converged Ethernet）是 InfiniBand 的主要替代方案。带宽可以做到类似，但延迟通常高于 InfiniBand，且需要额外的拥塞控制（PFC/ECN）来模拟 IB 的无损特性。

NVIDIA 的 Spectrum-X 平台就是基于以太网的 AI 网络方案，面向不想用 InfiniBand 的客户。

六、带宽层级

带宽 (GB/s, 双向) 技术 用途
─────────────────────────────────────────────────────────
8,000 HBM3e (B200) GPU 内部显存带宽
3,600 NVLink 6.0 (Rubin) 机柜内 GPU-GPU (未来)
1,800 NVLink 5.0 (B200) 机柜内 GPU-GPU
 900 NVLink 4.0 (H100) 节点内 GPU-GPU
 800 8×NDR IB (DGX H100) 节点间总带宽
 600 NVLink 3.0 (A100) 节点内 GPU-GPU
 200 XDR IB (单卡) 节点间单卡带宽
 128 PCIe 5.0 x16 CPU-GPU
 100 NDR IB (单卡) 节点间单卡带宽
 50 HDR IB (单卡) 节点间单卡带宽
 25 EDR IB (单卡) 节点间单卡带宽 / 100GbE

数量级关系:
 HBM 带宽 > NVLink > 多卡 IB 总带宽 > 单卡 IB >> PCIe

七、延迟层级

带宽之外，延迟同样重要：

操作 典型延迟
GPU 内部 SRAM (共享内存) 访问 ~几十 ns
GPU HBM 访问 ~100-300 ns
NVLink GPU-GPU ~1-2 μs
PCIe GPU-CPU ~2-5 μs
InfiniBand RDMA (节点间) ~1-3 μs
RoCE (节点间) ~2-10 μs
TCP/IP 以太网 (节点间) ~10-100 μs

LLM 学习日志 #1 并行方案和通信模式

Thu, 02 Apr 2026 17:20:54 +0800

在开始下一个阶段的开发之前，要先打好理论地基。虽然这部分之前零零散散的学过很多，但在通常的技术博客中通常只讲解每一种并行模式的技术原理和思想，将其和通信模式逐一详细联系起来并且做通信次数、数据量整合，甚至和反向传播的训练过程整合起来的并不多。

这篇文章打算稍微弥补一下这个问题，把所有零散的、常见或者不常见的并行切分方案，并且逐一评估其通信的需求、消耗、适用情况，在我的单人用户场景下，是否值得实现、实现难度如何。下面是总结。如果后面遇到了（甚至产生了）新的并行方案，也会一并记录在这里。

此外，在后面一篇博客当中，也会介绍一下每一种通信模式，以及通信模式对应的集群拓扑结构和相应常见的具体实现有哪些。

1、Tensor Parallelism（TP）— 层内切分

这个应该是最经典的并行方式了。它的切分方式就是把每一层的权重矩阵沿着head或者hidden维度切到多卡上，每卡算一部分，合并结果。不同的卡持有不同的head、或者FFN的权重矩阵的不同部分。

特别值得一提的是，Megatron-LM中一种经典的做法是，FFN中有两个线性层，可以把第一个线性层（W_gate和W_up）按列切、第二个按行切。这样的好处是，在中间的计算过程中，中间结果不需要通信。

它的流程如下：先用被切分的，只有自己这部分head的W_Q、W_K、W_V矩阵乘以输入得到自己head的Q、K、V，然后用Attention算出自己head的输出，然后把自己的head送入按照同样模式切分的、自己本地的相应的W_o部分做乘法，得到一个部分和，然后进行all reduce。值得注意的是，由于切分是相同的，因此绕过了单卡情况下，对head的全局cat操作因此无需通信，只需要cat自己本地的head就可以了。all reduce之后，送入自己本地的第一个FFN权重，然后逐元素激活，送入第二个权重，算出本地部分和，然后进行第二次all reduce。

通信模式：All Reduce

通信次数：每个Transformer Block里有两次All-Reduce（Attention之后一次，FFN之后一次）。总共次数为2*模型层数。

通信量：每一层中，涉及到的需要通信的，相关的每层数据量 = batch_size × seq_len × hidden_size × dtype_bytes。在整个模型中，总数据量 = 每层数据量 × 2（每层两次） × layer_num（层数）。对于ring All-Reduce的情况，总通信量为总数据量 × TP 度 = 总数据量 × 2 × (n-1)/n。

涉及的模块：Attention、FFN

优点： 1、可以达到很细的切分粒度，因为不同Head、FFN权重都可以切分，每张卡的显存占用可以按 TP 度线性减少。 2、实现很成熟，是最早的并行模式之一，有现实的工业级别落地。 3、天然是完全负载均衡的。

缺点： 1、通信频率很高，对卡间带宽要求很高，产生的网络延迟在非NVlink的机内和机间基本不可接受 2、切分过细的矩阵可能导致GEMM的算术强度降低，从而降低效率

训练中反向传播的情况

Megatron-LM给出了一个很漂亮的结论：对于所有前向的All-Reduce情况，在反向中的算子为Identity（不作数学操作，仅把同一个输入x复制到每张卡上）；对于所有的前向的Identity情况，在反向中的算子为。因此，在反向传播中，通信产生的次数和前向传播一模一样、数据量也相同，但出现的位置和前向传播相反。前向传播在结尾处做all reduce，而反向传播在开始处（当然同时是反向传播自己的结尾处）做all reduce。

这其实很符合直觉，简单推导即可得到，在此按下不表。

需要注意的是，反向传播的一个stage，在一般论文里大约占据两个slot，也就是两个时间步。换言之，被建模为前向传播的时间的两倍。

2、Pipeline Parallelism（PP）— 层间切分

Layer 分离 — 切分不同的Transformer Block

最简单直接的并行模式。也就是直接把模型的不同层分配到不同卡上。例如，对于一个32层模型进行对4卡的PP，那么每一张卡恰好持有 8层。数据从第一卡计算完成后送入下一张卡，以此类推直到流向最后一卡，像流水线一样。

通信模式：点对点（Point-to-Point）

通信次数：整个前向传播过程中总通信次数N-1次，N为卡数。

通信量：每次切换卡，单次卡间通信量 = batch_size × seq_len × hidden_size × dtype_bytes，一次前向传播的总通信量 = num_gpus × 单次卡间通信量。

涉及的模块：不涉及具体模块，按卡的stage进行划分，粒度为layers。

优点： 1、通信量小、频率低，可以其通信量和频率要求允许跨节点进行通信。 2、实现简单、计算规整，可以保证层内计算和单卡完全相同，不产生任何浪费和算术强度降低。 3、反向传播计算流和单卡情况相同，不需要对反向传播的组织模式进行特别的处理。

缺点： 1、存在流水线空泡（Pipeline bubble）问题，流水线的启动。 2、对于训练场景有额外的问题：反向传播需要反向流水线，交叉进行在简单的调度策略下会造成更多空泡，需要micro-batch、1F1B调度等等更复杂的策略。 3、对于层数不能被卡数整除的场景，会产生负载不均衡。此外，首尾层的embedding/lm_head可能产生lagging。

训练中反向传播的情况

Pipeline Parallelism的反向传播在数学上和单卡完全一致，每个stage各自做标准的反向传播就行。但难点在于调度：怎样安排各stage的前向和反向的执行顺序，使得GPU尽量少空闲。

需要了解的调度策略如下：

1、GPipe 朴素方案。先把所有 micro-batch 的前向做完，再统一做反向。会产生很多空泡，而且需要保存大量的micro-batch的激活值。值得注意的是：重计算想法（只存储部分中间激活值以节省显存）也是在这部分研究当中一同产生的。

2、1F1B（One Forward One Backward）在第一个forward做完之后立刻开始第一个backward的反向过程，和后面的forward过程穿插。通过组织调度数学上可以保证每个stage在任意时刻最多只需缓存PP个micro-batch的激活值。它的核心是：在稳定阶段，每个

3、Interleaved 1F1B 每个GPU不再持有连续的若干层而是持有总数相同但是在模型前向传播过程中不连续的层。这等价于增加了stage的数量，只不过可能有多个stage在同一张卡上。比如PP=4的情况下，8层模型、interleave=2；那么第一个stage（gpu）持有0、4层，第二个stage持有1、5层，第三个stage持有2、6层，第四个stage持有3、7层。它的好处是可以通过很简单的方式增加流水线深度，从而允许更少的流水线空泡，坏处则很明显：大大增加了需要的通信量，进而增加了设备需求和延迟。

4、Zero-Bubble Pipeline 这是一个很巧妙的想法。它把反向传播的过程中的成分做了更细粒度的划分。具体来说，反向传播其实包含两个独立的计算：

B（backward for input）：计算 ∂L/∂x，即对输入激活的梯度
W（backward for weight）：计算 ∂L/∂W，即对权重的梯度

前者被前一个stage依赖的，前一个stage只有获得了B才能开始自己的反向传播计算过程。而后一个实际上是悬挂在依赖链上的分支，它依赖后一个stage的结果但本身并不被其他人依赖。W的结果只是累加到本stage的grad buffer里，等最后optimizer.step() 时用。它什么时候算都行，只要在optimizer.step()之前算完就行。那么一个自然的想法就是：能否先算有依赖的B，让后面的先开始，在有空的时候自己慢慢算W？答案是可以的。

通过手工调度，或者基于**ILP（整数线性规划）**求解的自动调度，可以最小化bubble的数量，甚至于接近0bubble的比例。当然，它实际上会引入一个问题：由于W延后，实际上一个step的总Latency变长了。但通过约束，可以完全保证跨step重叠至多在一个step以内，对训练收敛几乎没有可观测的影响。

Prefill-Decode 分离（PD Disaggregation）— 切分请求的不同阶段

Prefill 和 Decode 的计算特性完全不同：

Prefill: 大矩阵乘法，compute-bound，高算术强度，适合高算力卡
Decode: 矩阵-向量乘，memory-bound，低算术强度，适合高带宽卡

混在一起跑会互相干扰。Decode的延迟敏感请求被Prefill的大计算堵住，Prefill的高吞吐被Decode的频繁小请求打断。

更进一步的说，两者的基本的算术强度就不同：Decode的算术强度一般可以估算为，FFN部分为engine中的batch数量的两倍（fp16或bf16意义下。fp8则乘以2）、Attention部分固定为1。实测数据显示，总体强度一般是40-90 FLOPs/byte左右，远低于H100 BF16的ridge point 295。而Prefill的算术强度随着seq_len同步线性增长，一般远远超过任何gpu的ridge point。

Decode 阶段

线性层（FFN、QKV 投影、Output 投影）

Decode 每步只生成 1 个 token，所以是矩阵-向量乘法（GEMV）：

权重 W: [d, d]
输入 x: [batch_size, d] （每个请求只有 1 个 token）

FLOPs: 2 × batch_size × d²
数据搬运: d² × bytes_per_param（权重）+ batch_size × d × bytes（输入，通常可忽略）

权重远大于输入，所以：

算术强度 ≈ 2 × batch_size / bytes_per_param

精度	batch=1	batch=8	batch=32	batch=128
BF16 (2B)	1	8	32	128
FP8 (1B)	2	16	64	256

Attention（decode 的 score 计算）

每个请求要用 1 个 query token 去和自己的整个 KV cache 做 attention：

Q: [batch_size, 1, d]
K cache: [batch_size, seq_len, d] ← 每个请求有独立的 KV cache

FLOPs: 2 × batch_size × seq_len × d
数据搬运: batch_size × seq_len × d × bytes（KV cache，每个请求各自的）

注意 KV cache 不能跨请求共享（每个请求的上下文不同），所以 batch 增大时数据搬运也等比增大：

算术强度 ≈ 2 / bytes_per_element ≈ 1（BF16）

不随 batch size 增长。

而Prefill的算术强度可以如此计算：

Prefill 阶段

线性层

Prefill一次处理整个 prompt，所以是矩阵-矩阵乘法（GEMM）：

权重 W: [d, d]
输入 X: [batch_size × seq_len, d]

FLOPs: 2 × batch_size × seq_len × d²
数据搬运: d² × bytes_per_param（权重）
 + batch_size × seq_len × d × bytes（输入激活）

当 batch_size × seq_len 足够大时，权重搬运的成本被大量计算分摊掉：

算术强度 ≈ 2 × batch_size × seq_len / bytes_per_param

精度	seq=1024, batch=1	seq=4096, batch=1	seq=4096, batch=4
BF16 (2B)	512	2048	8192
FP8 (1B)	1024	4096	16384

这些数字远远超过任何 GPU 的 ridge point。所以 Prefill 的线性层基本都是 compute-bound。

Attention（prefill 的 self-attention）

Prefill 中 Q、K、V 的长度都是 seq_len：

Q: [batch_size, seq_len, d]
K: [batch_size, seq_len, d]

Q×K^T 的 FLOPs: 2 × batch_size × seq_len² × d_head × num_heads
数据搬运: batch_size × seq_len × d × bytes × 3（Q, K, V）

算术强度大约是：

算术强度 ≈ seq_len / (3 × bytes_per_element)

这些差异是如此的明显，以至于这个disaggreation是最先被提出的。因此，分离之后各自可以针对性地优化硬件配置和调度策略。

Attention-FFN分离（AF Disaggregation）— 切分Transformer的不同阶段

AF分离是一个相对来说更新的分离模式，而其和MoE的诞生关系密切，是一个很细粒度的分离。这是因为Expert的“专家”指的就是不同FFN的专家，它们共享Attention的所有参数，但各自具有各自的FFN层和相应参数。更进一步的说，两者占据显存的东西并不相同，Attention部分主要是因为KV Cache占据了大量显存，KV cache在不同卡之间的来回拷贝可能造成Attention性能的严重下降；而FFN中，权重占据了显存的绝大多数（甚至可以达到80%以上）。对于一个MoE来说，显存的膨胀同时也允许（或者说需求了）更细粒度的分离模式。此外，MoE的通信模式本身就需要大量的、不规则的、不可预知的点对点通信，Attention→FFN的传输可以和All-to-All的dispatch合并成一步，这意味着AF分离的边际代价相较于dense模型的情况减少了。反正都是无法避免要付出的，不如一口气全做完。

总的来说，Attention更倾向于小TP度（保持GEMM效率，减少KV cache副本和通信，而且Attention并没有非常庞大的参数量导致的需求），而FFN/Expert想要大EP度（容纳更多expert）。可以说MoE天然适合AF分离——既在动机上需要它（Attention和Expert的并行度不匹配），又在代价上容忍它（All-to-All吸收了额外通信）。

3、Data Parallelism（DP）— 请求间切分

这应该是最朴素的并行了，朴素到不需要怎么讲。它的复杂度不在于推理（推理就完全是不同的实例在独自工作，无通信），而是在于训练的冗余参数存储的优化，这里诞生了许多大模型训练重要的早期成果——Google Brain、ImageNet、GPT-2、BERT，一直到单个模型在单卡上彻底放不下为止。

通信模式：推理无（训练则是All Reduce（经典训练DP）/All-Gather（ZeRO-1）/Reduce-Scatter（ZeRO-2）/两者（ZeRO-3））

通信次数：推理0，训练则是每个训练步 2 次集合通信（1 次 Reduce-Scatter + 1 次 All-Gather）（ZeRO-1）/ L+1 次集合通信（L 次 Reduce-Scatter（反向传播中逐层做）+ 1 次 All-Gather（更新后同步权重））（ZeRO-2）/ 3L 次集合通信（前向 1 次 All-Gather + 反向 1 次 All-Gather + 反向 1 次 Reduce-Scatter）（ZeRO-3）。

通信量：推理0，训练则是2 × Φ（可训练参数量） × dtype_bytes（经典DP和ZeRO-1和ZeRO-2）/ 3 × Φ（可训练参数量） × dtype_bytes（ZeRO-3）。

涉及的模块：推理无，训练涉及到优化器状态（ZeRO-1）、梯度状态（ZeRO-2）和权重（ZeRO-3）。

优点： 1、无敌简单和容易理解。 2、推理完全无通信开销。 3、天然是完全负载均衡的。

缺点： 1、对大模型的一个卡装不下的情况毫无帮助。 2、训练的反向传播过程依然需要通信。

训练中反向传播的情况

这是ZeRO-1、2、3和FSDP大显身手的地方。

ZeRO-1：切分优化器状态

切分方式：优化器状态（Adam 的 m 和 v）按参数均分到 N 张卡。权重和梯度每卡完整保存。

通信模式：Reduce-Scatter + All-Gather

通信过程：

反向传播结束后:
 1. Reduce-Scatter 梯度
 每卡得到自己负责的那 1/N 参数的聚合梯度
 通信量: Φ × dtype_bytes（每卡发出完整梯度，收到 1/N）

 2. 每卡用聚合梯度更新自己那 1/N 的参数（本地操作，无通信）

 3. All-Gather 更新后的权重
 每卡把自己更新的 1/N 权重广播出去，收集到完整权重
 通信量: Φ × dtype_bytes

通信次数：每个训练步 2 次集合通信（1 次 Reduce-Scatter + 1 次 All-Gather）

每步总通信量：2Φ × dtype_bytes（和经典 DP 的 All-Reduce 相同）

ZeRO-2：切分优化器状态 + 梯度

切分方式：优化器状态和梯度都按参数均分到 N 张卡。权重每卡完整保存。

通信模式：Reduce-Scatter + All-Gather

通信过程：

反向传播中（逐层）:
 1. 每层算完梯度后立刻做 Reduce-Scatter
 每卡只保留自己负责的 1/N 梯度，其余丢弃
 → 梯度显存从 2Φ 降到 2Φ/N

反向传播全部结束后:
 2. 每卡更新自己负责的 1/N 参数（本地操作）

 3. All-Gather 更新后的权重
 通信量: Φ × dtype_bytes

通信次数：每个训练步 = L 次 Reduce-Scatter（反向传播中逐层做）+ 1 次 All-Gather（更新后同步权重）

每步总通信量：2Φ × dtype_bytes（仍然和经典 DP 相同）

ZeRO-3 / FSDP：切分优化器状态 + 梯度 + 权重

切分方式：优化器状态、梯度、权重全部按参数均分到 N 张卡。每卡只存 1/N。

通信模式：All-Gather + Reduce-Scatter（前向和反向每层都要做）

通信过程：

前向传播（逐层）:
 对于第 i 层:
 1. All-Gather 收集完整权重 通信量: Φ_i × dtype_bytes
 2. 用完整权重做前向计算
 3. 丢弃非本卡的权重（释放显存）

反向传播（逐层，从后向前）:
 对于第 i 层:
 4. All-Gather 再次收集完整权重 通信量: Φ_i × dtype_bytes
 5. 计算梯度
 6. 丢弃非本卡的权重
 7. Reduce-Scatter 梯度 通信量: Φ_i × dtype_bytes
 每卡只保留自己负责的 1/N

全部反向结束后:
 8. 每卡更新自己的 1/N 参数（本地操作，无通信）
 （不需要额外 All-Gather 同步权重，因为下一步的前向会自动收集）

通信次数：每层 3 次集合通信（前向 1 次 All-Gather + 反向 1 次 All-Gather + 反向 1 次 Reduce-Scatter），总计 3L 次

每步总通信量：

前向 All-Gather: Φ × dtype_bytes （所有层加起来）
反向 All-Gather: Φ × dtype_bytes
反向 Reduce-Scatter: Φ × dtype_bytes
──────────────────────────────────────
总计: 3Φ × dtype_bytes

比经典 DP / ZeRO-1 / ZeRO-2 的 2Φ 多了 50%。多出来的就是前向传播中那一次 All-Gather——因为权重也切了，前向时也要收集。

4、Expert Parallelism（EP）— 专家参数间切分

这个同样是相对来说很容易理解的一种并行模式，也是MoE的必然选择，可以说这个并行模式和MoE架构是严格相互绑定的。这么做的主要原因有两个：第一，Expert的总量非常大，导致相同量级的MoE模型，计算量低得多，但所有Expert加起来的参数量远大于dense模型，导致模型根本不可能放在同一张卡上；第二，这些expert之间是独立的，不同token去不同的expert，中间过程不需要通信或者拼接，天然适合分布到不同卡上。

通信模式：All-to-All-v（数据量不规则）

通信次数：每个 MoE 层: 2 次 All-to-All（dispatch + combine），非 MoE 层（如 Attention）: 按 TP 的通信模式（All-Reduce）。

通信量：不可精确预测。每次All-to-All的通信量取决于token的分发模式，进一步地取决于router。理想均匀情况下:

每个 GPU 持有 batch_size × seq_len 个 token
每个 token 选 top-k 个 expert
总共需要分发的 token 数: batch_size × seq_len × top_k
其中留在本卡的比例: EP_degree 分之一
需要发送的: batch_size × seq_len × top_k × (1 - 1/EP) × hidden_size × dtype_bytes 

每个 MoE 层: 上述 × 2（dispatch + combine）

总通信量: 上述 × MoE层数。

但实际情况往往（高度）不均匀，热门 expert 收到的 token 多，冷门 expert 收到的少。

涉及的模块：FFN。

优点： 1、可以在不让计算量增加到难以承受程度的情况下，获得接近大参数量模型的性能。 2、切分粒度可以非常高。

缺点： 1、几乎总是一定有明显的负载不均衡，且无法在离线情况下静态的提前规划调度。 2、对显存的需求非常非常大。 3、训练比dense模型更加困难。

负载不均衡是MoE最难以处理的问题。Router是动态的，每个batch的token分发模式都不同。某些expert可能突然变"热门"，收到大量 token，而其他expert空闲。这会造成GPU的lagging情况，从而拖慢整个集群的性能。EP也是唯一一个通信模式动态变化的并行策略。这使得它的优化比其他并行困难得多，因为没法提前知道每一步的通信模式是什么样的。对它的调度策略比起前面的精密的线性规划模式，可能更接近于时序预测、控制论、排队论、资源调度等等范畴。可以想见，关于MoE模型的网络设计、集群设计、调度策略设计将会比前面的模式更有可以做的地方，可能也需要更长的时间才能在工业界收敛到一个相对来说足够优秀的统一策略，也可能根本不收敛。

训练中反向传播的情况

EP的训练问题并不在于技术上的问题——在DP、PP、TP等模式里已经基本上完全解决了。它面对的是专家坍缩和辅助负载均衡损失等等问题。它指的是一种Expert之间的马太效应：Router是可学习的，它在训练过程中会"偏心"——逐渐把越来越多的token路由给少数几个 expert，其他expert收到的token越来越少，最终几乎不参与计算。这是一个正反馈循环：某个expert因为偶然收到更多token，得到更多训练，变得更强，router就更倾向于选它，它就收到更多token……最终模型退化成一个dense模型，MoE的容量优势完全消失。更详细的介绍超出了这篇博客的范畴，因此就不多讲了。

此外，值得一提的是，Router的top-k选择是一个离散操作（不可微），但我们需要梯度回传。常用的做法是：

对被选中的expert，正常回传梯度
对没被选中的expert，梯度为零（它们根本没参与计算）
Router的梯度通过加权系数（gating weight）回传——这些系数是连续的、可微的

这意味着router的训练信号只来自被选中的expert，它对未选中expert的"好坏"一无所知。这也是expert collapse的一个根源。

5、Sequence Parallelism （SP） — 序列维度切分

SP是基于前文提到的TP的一种改进（同时也必须在TP的基础上做，无法单独存在。此外，它的名字实在是容易引起误解），它处理的是非FFN和非Attention的、没有被TP切割的地方，如LayerNorm的激活值、Dropout的激活值、残差连接的激活值等等，当然也是一个训练专门使用的优化手段。这是因为，激活值存储这个问题在推理里根本不存在。SP优美的地方在于，它不增加通信量，而是通过拆解通信原语的方式无代价的获得收益——也许可以被称为一种通信的Kernel Fusion？

回顾TP用到的通信模式：每个Attention/FFN模块结束时做一次All-Reduce。而在底层通信原语的逻辑上，All-Reduce其实一般是拆成两步实现的：

All-Reduce = Reduce-Scatter + All-Gather

SP 把这个拆分利用起来了：

不用 SP 时:
 TP 模块结束 → All-Reduce → 每卡得到完整结果 → LayerNorm(完整)

用 SP 时:
 TP 模块结束 → Reduce-Scatter → 每卡得到 1/N 的结果 → LayerNorm(1/N)
 → 下一个 TP 模块开始前 → All-Gather → 收集完整输入 → TP 计算

通信模式：Reduce-Scatter + All-Gather（替代 TP 原本的 All-Reduce）

通信次数：和 TP 相同，每个 Transformer Block 4 次集合通信（前向 2 次 + 反向 2 次），只是把 All-Reduce 拆成了 RS + AG

总通信量：Reduce-Scatter + All-Gather = All-Reduce。通过把All-Reduce拆分和还原，做到了在不增加额外开销的情况下减少大量激活值存储。

6、Context Parallelism （CP） — 长序列 Attention 切分

这是最近比较新的一种并行模式。CP应该是目前主流并行策略中最新加入的一维。Llama 3的训练就用了4D并行：TP + CP + PP + DP，CP 是第四维。

CP和Ring Attention技术密切相关。它是把输入序列沿token维度切分到多张卡上。每卡只持有序列的一段，以此实现利用并行性来加速的效果。

对于非Attention的操作（FFN、LayerNorm等），token之间没有交互，每卡独立计算自己的部分即可，完全不需要通信。唯一需要通信的是 Attention，因为每个token的Query需要和所有token的Key、Value交互。

Meta 在Attention模块中实现了两种CP变体，通常被称为Ring Attention。

Pass-KV（经典Ring Attention）

每卡保留自己的 Q 不动
K、V 块在 GPU 之间沿环传递
每轮：用本地 Q 和收到的 KV 块计算部分 Attention

标准Ring Attention，适合大多数场景。

Pass-Q（反向传递）

每卡保留自己的 K、V 不动
Q 块在 GPU 之间沿环传递
每轮：用收到的 Q 和本地 KV 计算部分 Attention

Pass-Q在某些特定场景下更优。当KV cache命中率低于5%时，Pass-KV更优。

Pico-vLLM 开发日志 #9 阶段性总结和开发过程的错误汇总

Tue, 31 Mar 2026 22:10:08 +0800

这篇简单的总结一下我们目前达成的里程碑、我们已经实现的功能和feature，同时也总结一下遇到的主要问题，提醒自己、也帮助后来人。

里程碑记录

注：有AI辅助排版。

开发时间线

从第一行代码到单卡收尾，总计11 天（3月21日 — 3月31日）。

Day 1–2（3/21–3/22）：从零搭建

从sampler.py开始，逐层手写Qwen2.5-1.5B的所有算子：RMSNorm、RoPE（复数形式）、GQA Attention、SwiGLU FFN，每个模块写完立刻和 HuggingFace逐tensor做torch.allclose数值对齐。完成权重加载、Engine、Sampler，在 GPU 上跑通第一次对话输出。

v0.1-naive-no-cache — 21tok/s，无 KV Cache的baseline。

Day 3（3/23）：KV Cache + Continuous Batching

实现NaiveKVCache，验证KV Cache在长序列下的加速效果（seq_len=3200时从1.8 tok/s提升到15.3tok/s，8.5x）。紧接着实现Scheduler和Continuous Batching，请求动态进出batch。

v0.2-naive-kvcache — 第一次 profiling，建立各算子耗时的数据基线。

v0.3-continuous-batching — Scheduler + Engine多请求调度跑通。

Day 4–5（3/24–3/25）：PagedAttention 数据结构 + 模型重构

实现BlockManager（全局block pool，分配/回收）和PagedKVCache（block_table + 逻辑→物理映射）。GQA Attention完全重写，拆分prefill和decode两条路径。RoPE 从复数形式重构为查表式实数实现。端到端数值验证通过。

v0.4-paged-attention — PagedAttention gather实现，性能和NaiveKVCache持平（这是因为GEMM主导，gather开销被掩盖）。第二次profiling。

Day 5–7（3/25–3/27）：Triton Kernel 密集期

手写两个Triton kernel：store_kvcache_kernel（2D grid，token × head）和Decode_Paged_GQAAttention_Kernel（online softmax + block table paged access）。完成QKV Fused Projection和SwiGLU gate_up Fused。将model.forward重构为纯tensor接口，所有Python对象操作移出forward，为CUDA Graph和torch.compile扫清障碍。

v0.5-triton-paged-attention — 23.2 tok/s。第三次profiling（nsys），首次量化 CPU launch overhead（851 kernels × 15μs = 13.5ms/步）。

v0.6-slot-mapping-store-kvcache — 28.4 tok/s，D2H memcpy 从560次降至0次。

Day 8–9（3/28–3/29）：CUDA Graph

先做实验验证CUDA Graph的收益预期（30 kernel串联实测6.62x加速），然后实现完整的静态buffer预分配 + capture/replay机制。Decode路径走CUDA Graph，Prefill保持eager。

v0.7-cuda-graph — 87.3 tok/s，4.2x加速。第四次profiling。

Day 10–11（3/30–3/31）：Triton Kernel 收尾 + 调优

手写Fused RoPE + Paged KV Cache Store kernel（+5 tok/s）、Fused SwiGLU kernel（+2 tok/s）、Triton RMSNorm（+2~3 tok/s）。尝试Fused Add RMSNorm，确认为负优化后回退。CPU侧block_table缓存优化。Engine收尾整合。

v0.8-97tok-cuda-graph — 96.94 tok/s，超过vLLM在同硬件上的95tok/s。

优化阶段 延迟 tok/s vs 起点
─────────────────────────────────────────────────────────
v0.1 Naive（无 KV cache） 51ms 21~1.5 1.0x
v0.2 NaiveKVCache 50ms 20 1.0x
v0.4 PagedAttention (gather) 48ms 20.8 1.0x
v0.5 Triton Paged Attention 43ms 23.2 1.1x
v0.6 + QKV/SwiGLU fused 35ms 28.4 1.35x
v0.7 + CUDA Graph 11.45ms 87.3 4.2x
v0.8 + Triton Kernel 收尾 10.32ms 96.94 4.6x

功能与特性实现

模型层

从零手写 Qwen2.5-1.5B 完整前向推理，不依赖 HuggingFace 的模型实现，仅使用其 Tokenizer 和权重文件。包括 RMSNorm、RoPE（查表式实数实现）、Grouped Query Attention（12 Q heads / 2 KV heads）、SwiGLU FFN，全部模块经过与 HuggingFace 参考输出的逐 tensor 数值对齐验证。支持 bfloat16 精度推理。

KV Cache 管理

实现了完整的 PagedAttention 内存管理体系。BlockManager 在启动时一次性预分配全局 KV Cache 物理内存池，按固定大小的 block（16 tokens）进行分配和回收，通过逻辑→物理的 block_table 映射实现非连续内存的透明访问，彻底消除显存碎片。每个请求持有独立的 PagedKVCache 对象，记录自身的 block_table 和序列长度，生命周期与请求绑定。

调度与 Batching

实现了 Continuous Batching 调度器，请求可以在任意时刻动态加入和退出 batch，不需要等待凑满固定 batch size。调度器维护 waiting、prefilling、decoding、finished 四个队列，每步自动完成状态流转。Prefill 和 Decode 走独立的执行路径，Prefill 逐请求处理（序列长度动态），Decode 支持多请求 batch 并行。

自写 Triton Kernel

共 5 个手写 Triton kernel：

store_kvcache_kernel：2D grid（token × head），通过 slot_mapping 将新计算的 K/V 直接写入非连续的物理 Cache 块，支持 stride 处理非连续内存布局。

Decode_Paged_GQAAttention_Kernel：Paged Decode Attention 的核心计算，在 kernel 内部通过 block_table 索引逐块读取分散存储的 KV，使用 online softmax 累积 attention 结果，天然处理 GQA 的 head 映射，零拷贝无需 gather。

fused_rope_and_cache_store_kernel：将 RoPE 旋转位置编码和 KV Cache 写入融合为单个 kernel，Q 在寄存器内完成旋转后写回连续显存，K 和 V 完成旋转后直接通过 slot_mapping 写入分散的物理 Cache 块，消除中间缓冲。

fused_silu_mul_kernel：将 SwiGLU 的 Sigmoid、门控乘法和 up 分支的逐元素乘法融合为单次显存读写，三次 I/O 降为一次。

_rmsnorm_kernel：将 RMSNorm 的 6 个独立 PyTorch 算子合并为单个 kernel，每行一个 program，全程 float32 计算后转回原始 dtype 写出。

算子融合与模型优化

QKV Fused Projection：三个独立的 q/k/v 矩阵乘法合并为单次 GEMM，权重在加载时拼接，输出通过 split 拆分。SwiGLU gate_up Fused：同理将 gate_proj 和 up_proj 合并。rotate_half 的 torch.cat 消除，改为 in-place 操作。model.forward 重构为纯 tensor 接口，不接受任何 Python 对象，所有 cache 操作留在 Engine 层，使 forward 对 torch.compile 和 CUDA Graph 完全透明。

CUDA Graph

Decode 路径完整接入 CUDA Graph。Engine 启动时预分配固定形状的静态 buffer（input_ids、slot_mapping、position_ids、block_table、context_lens），每步只原地更新 buffer 的值，不重新分配内存。CUDA Graph 一次 capture 后反复 replay，将每步数百次 kernel launch 压缩为单次提交。Prefill 路径保持 eager 执行（序列长度动态，不适合 capture）。提供 use_cuda_graph 开关，可随时回退 eager 模式调试。

采样

支持 Greedy、Temperature Sampling 和 Top-p (Nucleus) Sampling 三种采样策略，统一接口。

开发过程遇到的错误汇总

简单的总结了一下，整个单卡框架开发过程中遇到的主要的、难以排查的错误，我还记得的就这么几个——如果我不记得了，应该是没有折磨我太久，大家应该能解决。

注：有AI辅助排版。

类型一：KV Cache 写入位置错误

症状： 生成乱码，词语乱序重复，或反复覆盖同一位置。

出现了两次：

第一次（早期 PagedAttention）：

原因：context_lens 比实际少 1
decode 时写入新 token 后，读取范围没有包含刚写入的 token
新 token 的 query attend 不到自己

第二次（slot_mapping 重构后）：

原因：get_prefill_slot_mapping 里
cache_block_index 是 tensor，用 tensor 作为 dict key
Python dict 用 tensor 作 key 行为未定义，查到错误的物理 id
slot_mapping 返回 [0,1,2,3,4] 而不是真实物理地址
prefill 写入物理 block X，decode 读取物理 block Y
加 .item() 后修复，但这降低了性能，后续整体重构了

类型二：非连续内存导致 Triton Kernel 越界

症状： cudaErrorIllegalAddress，kernel 崩溃，或者生成完全不可理解（而不是简单反复重复）的乱码。

出现了两次：

第一次（paged_decode_attention）：

原因：kv_cache[layer_idx] 是大 tensor 的 view
Triton 拿到 data_ptr 后按连续内存计算偏移
实际 stride 和预期不符，导致越界访问
加 .contiguous() 后修复

第二次（store_kvcache，极其隐蔽）：

原因：apply_rope 之后 k/v 是非连续 tensor
view/reshape 假设连续内存
Triton kernel 按错误地址写入
根本修法：给 store_kvcache_kernel 传入 stride 参数
用 stride 计算偏移而不是假设连续

类型三：模型前向传播的过程实现在反复修改中出错

症状： 生成不完全的乱码，词语乱序重复，有意义但不完全有意义。

出现了一次，很隐蔽：

原因：q 相关的 forward 实现在不同路径里不一致，在 Decode 路径中遗漏了 RoPE 旋转状态
 - Eager 路径（forward）里，q 被 apply_rope 原地覆盖更新，传给 Attention 的是正确旋转后的 q
 - Graph 路径（forward_decode）引入了融合算子 fused_decode_rope_and_cache，它返回了新的 q_rot，但未被后续使用
 - 原始未旋转的 q 直接参与了注意力计算，导致模型在 Decode 时彻底丢失相对位置信息（RoPE 失效），引发无脑复读和乱码

出现位置：
 - GQAAttention.forward_decode 方法的末尾处
 - 算子调用错误地写成了 output = self._decode_attention(q, kv_cache_k, ...)
 - 实际应该传入旋转后的变量，改为 output = self._decode_attention(q_rot, kv_cache_k, ...)

类型四：Tensor Shape 错位

症状： RuntimeError: mat1 and mat2 shapes cannot be multiplied。

出现了多次，都是同一个模式：

原因：q 的 shape 在反复迭代修改后，在不同路径里不一致或者前后不一致。例如：
 - GQAAttention 内部 q 是 (B, seq, n_heads, head_dim)
 - paged_decode_attention wrapper 期望 (B, n_heads, 1, head_dim)
 - transpose 和 reshape 的顺序搞错导致维度错位

出现位置：
 - _decode_attention 里 q 传给 Triton 之前忘记 transpose
 - wrapper 里 unpack 顺序写成 B, _, N_HEAD, HEAD_DIM
 实际应该 B, N_HEAD, _, HEAD_DIM

类型五：seq_len 更新时机错误

症状： 生成内容在第 2 步之后开始乱码，第 1 步正确。生成的语句基本有正确的语义，但就是Decoding阶段开始的时候会出现重复、或者突然的错误。

原因：decode 每步用 context_lens = c.seq_len + 1
但 seq_len 在 forward 之前已经被 prepare_decode_step 
或其他逻辑提前更新，导致 context_lens 多 1 或少 1

另一次：prefill 后 _seq_len += len + 1（多加了 1）
第一个 decode step 的 position 偏移 1，之后全部错位

Pico-vLLM 开发日志 #8 更多Triton Kernel和性能优化

Tue, 31 Mar 2026 16:20:34 +0800

今天的主要工作是接入了更多的Triton Kernel，进一步完善代码和进行性能优化，并且把Decode的性能拉到了可以接受的程度。后面还有一篇阶段性的总结，主要讲一讲遇到的令我比较有印象的错误、呈现的形式和最终解决的方法。稍后更新。

小小的炫耀一下：

![Final Performance](Final Performance.png)

之前提到，vllm的性能是95.0tokens/s左右。已经和vllm完全持平甚至略微超过了！engine的适配和测试也已经整理完成，接下来就是准备重新开始学习、调研多卡的实现方法、设计多卡的扩展实现路线图了。

其他基于Triton Kernel的性能优化

除了之前的博客提到过的几个重要的Kernel或者不用Kernel的矩阵乘法合并优化，在完成 PagedAttention 的核心实现之后，nsys profiling 的数据显示GPU计算时间仍有优化空间。GEMM和GEMV在这个过程中已经占据了大约80%的GPU时间，这部分受限于硬件带宽，难以突破。但剩余的20%里藏着几个可以用Triton消除的低效模式：显存的重复读写和不必要的中间变量分配。通过剩下的零碎Kernel优化，仍然可以实现将近10%的性能提升。

1. Fused RoPE & Paged KV Cache Store (显存读写合并)

这是收益最大的一次融合，解决的是一个隐蔽的"显存回旋镖"问题。

在之前直接随手写的naive实现里，apply_rope会产生一系列临时tensor：先slice切出前后两半，再neg取负，最后cat拼回去。每一个操作都是一次独立的GPU kernel，意味着中间结果要反复写入显存再读出来。更糟糕的是，RoPE算完的K会先落入连续显存，然后store_kvcache再把它读一遍，按slot_mapping的地址写入分散的物理Cache块。同一份K的数据在没有然后额外处理需求的情况下被搬运了两次，造成冗余。

融合后的kernel把这个流程压缩成一次：从全局显存读入Q、K、V、cos、sin，在寄存器里完成RoPE旋转计算，全程不进行临时变量分配。计算完成后，Q写回连续显存供后续Attention使用；而寄存器里的K_rot和V则直接通过slot_mapping提供的物理地址，一步写入非连续的Paged Cache块，不再经过任何中间缓冲。

GQA的处理是这个kernel需要特别注意的一个设计细节。以框架适配的模型Qwen2.5-1.5B为例，Q有12个头，K和V只有2个头，Grid设计时需要让每个program明确知道自己处理的是Q头还是KV头，以及对应的KV头索引是哪个。通过kv_head_idx = head_idx // (N_HEAD // N_KV_HEAD)的映射，kernel内部自然地处理了这个不对齐问题，调用方不需要任何额外操作。

这个收益大约有5Tokens/s，也就是说足足5%。

2. Fused SwiGLU (MLP 层激活函数与乘法融合)

SwiGLU 是相对好发现和动机更直接的一个Kernel fusion操作，不过相对来说收益没有前面那个那么大。原始的F.silu(gate) * up`需要进行是三次显存IO：读gate、算SiLU、写临时结果；读临时结果和 up、做逐元素乘法、写输出。实际上这三步在计算上完全可以流水完成，没有任何数据依赖要求它们分开。

Triton kernel把gate和up同时加载到寄存器，在寄存器内连续执行Sigmoid、gate与Sigmoid结果的乘法（SiLU），以及与up的逐元素乘法，最后一次性写回全局显存。三次读写变成一次，访存带宽需求降低到原来的三分之一。

实现上把3D tensor（batch×seq×hidden）展平成1D视角来划分Block，这样kernel对任意的batch size和seq_len都能够天然兼容，不需要针对不同输入形状做特殊处理。

这个收益大约有2Tokens/s。

3. Triton RMSNorm (pow mean add sqrt和mul的算子融合)

这个算子fusion的目的是，把RMSNorm自身的多个零碎的PyTorch算子合并成一个Triton kernel。

原始的PyTorch实现 x.pow(2).mean(-1).add(eps).rsqrt().mul(x).mul(weight) 实际上在GPU端会拆成5~6个独立的elementwise kernel，pow一次、mean（reduction）一次、add一次、rsqrt一次，最后还有两次mul。每个kernel都有独立的 launch 开销，中间结果要落入全局显存再被下一个 kernel 读出来。对于28层×2个Norm=56次/步来说，这些launch虽然单次很小，累积起来也有可观的数量。

Triton kernel的做法很直接：每行（即每个token的hidden_size维度）分配一个program，一次性把整行1536个元素加载到寄存器，在寄存器内完成平方求和、除以N、加eps、rsqrt、乘权重的全部计算，最后一次写回。6个kernel变成1个，中间没有任何全局显存的临时读写。

这里要特别注意一个踩了的坑：必须用BLOCK_SIZE取hidden_size的下一个2的幂次（1536→2048），用mask屏蔽超出部分。这个实现实际上不是可选的二世必须的，来源于Triton对tl.arange的要求。具体来说，Triton要求tl.arange产生的长度必须是编译期常量且为2的幂。如果不是这样，会直接出错，在这里记录一下以提醒后来人。

此外，计算全程在float32下进行以避免bfloat16的精度溢出，写回时再转换为输入的原始dtype。

这个收益大约有2~3Tokens/s。

4. Fused Add Norm 失败

Fused Add Norm的融合是计划中的第4个优化，目标是在前面第三个的基础上更进一步，把残差加法和归一化也合并成一个kernel，从两次显存读写变成一次。实现上用Triton在Block内部完成方差计算和归一化，逻辑并不复杂。

动机很直接：在TransformerBlock里，残差连接和RMSNorm是紧挨着的两步操作。原始流程是先做x = x + attn_out（一次全局显存读写），再把结果送入RMSNorm（又一次读写：pow → mean → rsqrt → mul）。同一份数据在没有任何计算上的阻隔的情况下被搬运了两次。Triton kernel 的设计思路是把这两步压缩进一个 kernel：在寄存器里完成x + residual之后，就地计算方差、rsqrt 和缩放，x被in-place更新为加完残差的值，返回的是归一化后的结果，全程只读一次、写两次（一次更新 x，一次写 normed_out）。

但实际上测试下来不仅没有增加性能还减少了。是负收益，降低了大约2~3Tokens/s。尝试问问AI，给出的其中一种解释如下：

PyTorch 原生的 elementwise kernel 对这种小规模操作已经做了充分优化，走的是高度特化的 launch 路径；而 Triton kernel 有自己的一套 launch 机制，对于这种计算量极小、一个 warp 就能处理完的操作，Triton launch 本身的固定开销反而比省下来的显存带宽还大。更具体地说，在 Qwen2.5-1.5B 的 decode 阶段，hidden_size 只有 1536，batch_size 为 1，整个 RMSNorm 的输入只有 1536 个 float16 元素（3KB），连 L2 cache 都装得下，根本不存在"显存回旋镖"的问题——数据在 cache 里就地复用了，多读一次的代价几乎为零。

不过，我个人认为这应该就是launch overhead的问题。就是不知道pytorch的原生实现到底为什么能够更快，后台又对不同size做了哪些优化？因为按理来说，elementwise的操作应该完全没有什么特别的优化空间才对。别说Triton了，可能用CUDA几小时写一个naive版本也能达到接近可达上限的性能。那到底是为什么？这个问题可能需要后面有时间再对pytorch做进一步的了解，现在先记录下来🤔。

这个收益大约是-2~-3Tokens/s，确认为负优化后已回退。保留了Triton版本和相关注释没有删除，作为参考。

把所有这些全部加上，性能基本上就能够和vllm一样了。当然，这是单Batch的情况，多Batch由于Scheduler的调度策略非常不同，测试的意义也不是很大。

Pico-vLLM 开发日志 #7 CUDA Graph

Mon, 30 Mar 2026 23:50:26 +0800

在完成Paged Attention的完整开发之后，下一步就是考虑提升性能了。

先来看benchmark。vllm在5070 Laptop上可以达到单Batch情况下，95 Tokens/s的Decoding速度。考虑到Decode阶段是绝对的Memory Bound，而5070 Laptop的理论内存带宽上限是382GB/s，这差不多相当于78%的GPU利用率。与此同时，在我的框架中进行相同的测试在Attention部分用自己写的Triton Kernel替换、进行了QKV Projection Fusion等比较重要的合并之后，性能虽然略微提升，但是在5070 Laptop上也仍然只有大约25~27 Tokens/s。即使是在进行了.compile()的default模式的编译之后，性能提升也最多只能达到大约31 Tokens/s。Profiling得到的性能情况如下图：

中间这4倍（对于eager模式）或者3倍（对于compile模式）的差距去哪里了？答案是，时间都被GPU-CPU的传输延迟、Python端的处理和CUDA API的Launch Kernel的开销吃掉了。这一点在WSL的环境下体现的尤其明显：对于WSL环境，一次CUDA API的调用就需要足足15微秒，比windows的原生环境和linux的原生环境都要长很多。而在一次模型的forward的过程中，至少会有大几百次（对于我在进行Profiling的时候使用的forward计算流，每个完整的Decoding步会产生851次cudaLaunchKernel的调用——仅仅调用本身的开销就消耗了13.5ms，而所有的GPU Kernel的实际开销也才11ms左右。因此，如果想要进行Decoding的性能的进一步优化，优化GPU Kernel的边际收益是很低的——继续减少一个本来已经占比不大的部分所产生的时间开销，其加速比永远不可能超过这个部分本来占据的比例本身。关键在于把更大的这一块优化好。

什么是CUDA Graph

CUDA Graph是一种可以减少CPU Launch Kernel的overhead开销的一种方法。为此，我们需要先了解CPU Launch Kernel的overhead开销是如何产生的。

我们知道，GPU在以前的很多场景下可不被称为GPU——它的另一个名字是加速器，或者说加速卡（accelerator）。这个说法有它自己的道理：GPU不会自己给自己发任务，它总是通过等待调用，被调用，返回结果的模式，帮助CPU进行它难以短时间完成的任务，也就是加速。这就需要CPU主动发放指令，这就是Launch Kernel。每一次kernel launch，CPU 都要经历一套流程：准备参数、调用 CUDA Driver API、把launch命令推入CUDA stream。GPU那边则从stream里取命令执行。关键点是这个过程是串行且逐条下发的：CPU发一条，GPU 收一条。即使 GPU 端的kernel执行极快（比如几微秒），CPU端每次launch也都是不可变的固定开销（WSL下~15μs），这部分是省不掉的。可以立刻得到的一条推论是，Kernel的数量越多，CUDA Graph的收益越大，反之则没有什么用。

CUDA Graph为什么能加速Decoding

大模型的推理Forward过程其实恰好就是许多个零碎的小Kernel组成的——而不是大家想象中的少数几个很大的Kernel。这一方面是因为layer堆叠的很深，另一方面是Attention、SwiGLU、RoPE等等的计算需求各不相同，完全无法进行完整的Fuse。在eager模式下，这样的控制流需要CPU亲自介入反复发放——而这个发放的时间无法完全overlap，很多Kernel在GPU上工作的时间甚至远短于CPU的launch开销，更别提在这个过程中，可能还夹杂了Paged Attention的实现的非最优实现引发的CPU侧反复的Tensor创建、内存分配等等行为。如果不实现CUDA Graph，可能很多人都没有意识到自己在中间有哪些动态拷贝拉低了性能。因此，CUDA Graph的实现过程产生的副产品其实比CUDA Graph本身还要更多一些。

对于大模型本身来说，CUDA Graph造成的直接收益则是如下几条：

第一，CPU端从N次launch变成了1次，直接砍掉了N-1次的driver调用开销；

第二，GPU端由于提前知道了完整的执行计划，调度器可以做更激进的优化（比如提前准备下一个kernel的资源）；

第三，消除了CPU-GPU之间反复同步的 round-trip（不一定是Kernel本身造成的，正如我之前所说，可能有很多不经意之间的非最优实现）。

CUDA Graph的技术原理

一句话总结的话，它的本质其实是录制和重放。这有点类似于宏录制和宏重放的思想：把一套线性的、无条件分支发散的操作流程固定下来，然后在输入数据不同的情况下反复重做。这样，一整个过程本身只需要启动一次，后面GPU自己会知道怎么做。

CUDA Graph做的事情就是：把一系列kernel launch的指令序列，或者说launch的顺序、拓扑依赖关系，全部提前录制（capture）下来，形成一个有向无环图（DAG）。图中的节点是各个kernel（以及memcpy、memset等操作），边表示它们之间的依赖关系。录制完成后，这整张图可以作为一个整体一次性提交给 GPU 执行。CPU 端只需要一次 launch 调用（cudaGraphLaunch），GPU 端就会按照图中定义好的顺序和依赖关系依次执行所有节点。

如何使用CUDA Graph

这东西的使用其实很简单，不过它的限制也很大。它的使用方法真的就是字面意义上的“录制和重放”——先在开启录制的完整的跑一遍，然后调用replay相关的操作。从api角度来解释的话，流程如下：

cudaStreamBeginCapture → 正常调用各种 kernel → cudaStreamEndCapture → 得到一个 cudaGraph_t → cudaGraphInstantiate 生成可执行的 cudaGraphExec_t → 之后每次只需 cudaGraphLaunch即可。

具体而言，在pytorch的实现当中，可以参考这段代码的流程：

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


# ============================================================
# 准备Decode的静态数据
# ============================================================
cache.prepare_decode_step()
static_input_ids[0, 0] = next_token
static_slot_mapping[0] = cache.get_decode_slot()
static_position_ids[0, 0] = cache.seq_len
static_block_table[0, :cache.allocated_cache_block_num].copy_(cache.get_block_table())
static_context_lens[0] = cache.seq_len + 1

# 预热触发 compile 编译
for _ in range(3):
 with torch.no_grad():
 _ = model.forward_decode(
 static_input_ids,
 kv_cache_k=bm.gpu_kv_cache[0],
 kv_cache_v=bm.gpu_kv_cache[1],
 position_ids=static_position_ids,
 slot_mapping=static_slot_mapping,
 block_table=static_block_table,
 context_lens=static_context_lens,
 )
torch.cuda.synchronize()
# ============================================================
# 录制CUDA Graph
# ============================================================
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
 static_output = model.forward_decode(
 static_input_ids,
 kv_cache_k=bm.gpu_kv_cache[0],
 kv_cache_v=bm.gpu_kv_cache[1],
 position_ids=static_position_ids,
 slot_mapping=static_slot_mapping,
 block_table=static_block_table,
 context_lens=static_context_lens,
 )

next_token = static_output[0, -1].argmax()
print(tokenizer.decode([next_token.item()]))

# ============================================================
# Decode Loop
# ============================================================
PROFILING_TOKENS = 100 
torch.cuda.synchronize()
start_time = time.perf_counter()

for step in range(PROFILING_TOKENS):
 cache._seq_len += 1 
 cache.prepare_decode_step()

 #另一种原位写入方法，比前面的按索引写入更好
 static_input_ids.copy_(next_token, non_blocking=True)
 static_slot_mapping.fill_(cache.get_decode_slot())
 static_position_ids.fill_(cache._seq_len)
 static_context_lens.fill_(cache._seq_len + 1)

 bt = cache.get_block_table()
 static_block_table[0, :bt.shape[0]].copy_(bt, non_blocking=True)

 # 回放 Graph
 g.replay()
 # 贪心解码，在GPU侧完成以避免同步
 next_token = static_output[0, -1].argmax()

CUDA Graph的使用限制

所有“录制和重放”应该有的限制它都有。总结一条的话，它必须是线性的、静态的。

第一，所有的指针不能变动。换句话说，你不能使用任何“新内存分配”、“新张量创建”，所有这些操作本质上都是新分配内存空间、返回新指针，而CUDA Graph只接受静态的不变化的指针，那么当然也就只指向同一块空间、同一个Tensor。在不同的调用当中，能够做到的就是向固定的static的input里覆写不同的内容，以实现不同循环的不同输出。这也是为什么block_table在vllm实现里是定长的。

第二，控制流必须固定。 换句话说，录制时走过的分支就是回放时会走的分支，不能有data-dependent的 if/else。这意味着如果你的forward里有根据输入动态选择不同kernel路径的逻辑（比如根据seq_len是否超过某个阈值来决定用哪个attention实现），这些在 graph 内部是做不到的。所有的条件判断必须提到 graph 外面，在 CPU 端决定好之后选择对应的 graph 来 launch。

第三，kernel的launch配置（grid size、block size）和调用序列必须完全一致。 就像前面说的那样，一次录制记录下了851次kernel launch，回放时就真的只会严格按相同的顺序launch这851个kernel，grid/block参数也完全不变（作为静态变量，本质上，和前面的指针是一样的）。这直接意味着tensor shape不能变——因为绝大多数kernel的grid配置是根据tensor shape算出来的。对LLM inference来说，这一点的影响很大：decode阶段每一步kv cache的有效长度都在 +1，如果attention kernel的grid size依赖于seq_len，那严格来说每一步都需要不同的 graph。实际工程中的处理方式通常是用padding或者按seq_len区间预录制多张 graph。

第四，不能包含CPU-GPU同步操作。 比如cudaMemcpy（同步版本）、cudaDeviceSynchronize、或者任何会阻塞CPU等待GPU结果的调用，都不能出现在capture的过程中。因为capture阶段这些操作并不会真正执行，录进去会导致语义错误。同样，任何需要读回 GPU 数据到CPU再做决策的逻辑（比如early stopping，CPU介入检查EOS token）也必须移到graph外部。GPU必须表现的像CPU根本不存在一样来进行重放。此外，每一步检查EOS token这种操作会造成严重的性能影响（因为要进行同步和拷贝），一般是按照预定的长度，生成很多个（比如100个token），然后转交给CPU，检查有没有出现EOS token，是在哪里出现的，然后再进行进一步处理和截断。

下面这条是Claude老师补充的：

第五，不能在graph内部使用CUDA的动态并行（dynamic parallelism）或者跨stream的复杂同步。 Stream capture对stream的使用有严格约束，虽然支持在capture期间fork出子stream来表达并行性，但不能随意使用event做跨graph的同步。这部分我还没有研究过，不是很懂，先记下来。

总结一下的话，CUDA Graph本质上是把一段GPU执行流"快照"下来、机械性的、死板的重放。任何运行时才能确定的东西，包括新的内存地址、动态的控制流、变化的shape、CPU端的同步和决策等等，都和它不兼容，能改变的只有固定地址里的数据内容。

CUDA Graph的加速效果

非常好。直接上图⬇️：

这个直接达到了大约92%的vllm性能。后面通过进一步的Triton-Kernel的接入和fusion，应该完全可以做到93~95token/s，彻底追平vllm。

一点思考

在做这部分工作的时候，也在想Vera CPU这种东西会不会让Graph变得更不那么必要？

“Vera CPU 通过 NVLink-C2C 和 Rubin GPU 直接互联，提供 1.8 TB/s 的 coherent 带宽，构建了 CPU 和 GPU 之间的统一地址空间。”

也许以后。这就得等Rubin真的能摸到再说了。

Pico-vLLM 开发日志 #6 Triton Kernel和代码重构

Sun, 29 Mar 2026 21:43:05 +0800

一天没更新是因为花了点时间调试Triton Kernel。这部分的开发比我预期的要慢，主要还是对于pytorch的Tensor的操作不是非常熟悉，栽了几个比较大的坑，因此反复debug了好久。这里也记录一下，希望能以此提醒自己，也能够帮助后来者避坑。

这部分的思路其实前面已经完全阐述完成，主要是讲讲主要完成的工作和遇到的坑有哪些。

设计模式

这部分主要是为了后面的CUDA Graph的实现做准备。总的来说，指导思想只有一个：静态化所有核心前向的数据指针和形状，同时避免任何条件分支。这是因为CUDA Graph的本质是重放：将既定的指针形状、顺序、流程硬编码在GPU上，反复执行，从而避免CPU下发指令进行Kernel启动的开销。这部分目前还没有用上，但已经提前做了准备。

但众所周知，LLM的推理框架本身就需要两个截然不同的需求：Prefill和Decode。在vLLM等主流实践当中，一般都需要将两者完全分离，前者不使用CUDA Graph而后者使用。这么做的原因主要是Prefill的形状高度不固定：为一个比较短的prompt input去padding一个巨大无比的形状来确保绝对不越界以允许CUDA Graph明显是相当不方便的。这也是chunked Prefill产生的动机之一：分次做的话，很大程度上就形状固定了。如果将两者放在一起，is_prefill的条件判断本身也会让CUDA Graph因为分支发散而完全不成立。

这倒是让我想起来之前看到的一篇论文：《Medusa: Accelerating Serverless LLM Inference with Materialization》。这也是我这学期在上的课的老师陆游游老师团队做的一项工作。它的主要贡献是，直接把CUDA Graph的结构给解析了出来，然后通过物化和指针替换的方式动态的使得物化在主存或者SSD上的CUDA Graph不会因为重新加载之后cublas等闭源库的函数加载位置改变、数据指针的位置改变而失效，从而实现LLM推理服务从非加载状态的完全冷启动，同时还能保持启动的高效率，免去每次重新记录和重放的时间overhead消耗。考虑到CUDA Graph这东西能够在一个自研的框架上跑通本身就要花费巨量的心血，直接将其逆向的工作能够发在 ASPLOS 2025 上当然也是可以想象的。很推荐大家去读！

Triton Kernel的实现

目前显式写出的Triton Kernel有两个：一个是KV Cache 写入算子 (store_kvcache)，一个是分页解码注意力(paged_decode_attention)。它们分别负责：

1、store_kvcache在Prefill阶段和Decode根据block Manager分配的块表正确的写入相应的KV cache，Prefill阶段写入seq_len个而Decode阶段固定写入1个。

2、paged_decode_attention则是Decode阶段专用的Kernel，负责根据页表正确取出相应的KV cache，同时完成Flash Attention模式的不将中间矩阵物化的计算流程。

此外需要特别注意的是，虽然传入的是Tensor，不过实际上Triton内部的操作模式更接近C++：传入的Tensor实际上是隐式转换成带数据类型的指针来处理的。这方面要格外注意的是不同指针的数据类型：如果数据类型不同，那么指针的算术操作造成的物理地址偏移量实际上是完全不同的，而且写入的东西也很可能不匹配。如果不能确定自己到底在写什么的话，就多使用tl.print()然后只针对pid==0的情况进行打印来确认吧。

下面是源代码：

1、store_kvcache

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


@triton.jit
def store_kvcache_kernel(
 k_ptr, v_ptr, # 源 K/V: (total_tokens, n_kv_heads, HEAD_DIM)
 k_cache_ptr, v_cache_ptr, # 目标 Cache: (num_blocks, n_kv_heads, block_size, HEAD_DIM)
 slot_mapping_ptr, # (total_tokens,) 每个 token 的物理 slot 编号
 stride_k_token, stride_k_head, stride_k_dim,
 stride_v_token, stride_v_head, stride_v_dim,
 N_KV_HEADS: tl.constexpr, # ← constexpr
 BLOCK_SIZE: tl.constexpr, # ← constexpr
 HEAD_DIM: tl.constexpr,
):
 # 映射 2D Grid
 pid_token = tl.program_id(0)
 pid_head = tl.program_id(1)

 # 读取物理 slot
 slot = tl.load(slot_mapping_ptr + pid_token)
 block_id = slot // BLOCK_SIZE
 offset = slot % BLOCK_SIZE

 # 计算源数据 (K/V) 的一维内存偏移
 dim_offsets = tl.arange(0, HEAD_DIM)
 src_offset = pid_token * stride_k_token \
 + pid_head * stride_k_head \
 + dim_offsets * stride_k_dim

 k_vec = tl.load(k_ptr + src_offset)
 v_src = pid_token * stride_v_token \
 + pid_head * stride_v_head \
 + dim_offsets * stride_v_dim
 v_vec = tl.load(v_ptr + v_src)
 
 # 计算目标数据 (KV Cache) 的一维内存偏移
 dst_offset = (block_id * N_KV_HEADS * BLOCK_SIZE * HEAD_DIM) + \
 (pid_head * BLOCK_SIZE * HEAD_DIM) + \
 (offset * HEAD_DIM) + \
 dim_offsets

 # 写入 Cache
 tl.store(k_cache_ptr + dst_offset, k_vec)
 tl.store(v_cache_ptr + dst_offset, v_vec)

这个Kernel看上去非常简单，但是在优化的时候有一个大坑，后面讲。聪明的读者可能看到 stride 的时候就想到了吧。

2、paged_decode_attention

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


@triton.jit
def Decode_Paged_GQAAttention_Kernel(
 q , # (B, n_heads, 1, head_dim) query，decode 每步只有 1 个 token
 k_cache , # (num_blocks, n_kv_heads, block_size, head_dim) 全局 K cache
 v_cache , # (num_blocks, n_kv_heads, block_size, head_dim) 全局 V cache
 block_table , # (B, MAX_BLOCKS_PER_SEQ) int32，每个请求的物理块 id，-1 表示未分配
 context_lens , # (B,) int32，每个请求当前的有效 token 数
 scale , # 1.0 / sqrt(head_dim)
 out,
 
 # Meta-parameters
 # 元参数
 MAX_BLOCKS_PER_SEQ: tl.constexpr, # 启动时固定，不是运行时变量
 BLOCK_SIZE: tl.constexpr, #
 HEAD_DIM: tl.constexpr, #
 N_KV_HEAD: tl.constexpr,
 N_HEAD: tl.constexpr,
 ): # (B, n_heads, 1, head_dim)
 # grid = (B, n_heads)
 # 每个 program 处理一个 (batch, head) 对
 # program_id[0] = batch_idx
 # program_id[1] = head_idx
 
 pid_batch = tl.program_id(axis=0)
 pid_head = tl.program_id(axis=1)
 kv_head_idx = (pid_head // (N_HEAD // N_KV_HEAD))

 m = float('-inf')
 l = float(0)
 o = tl.zeros((HEAD_DIM, ), dtype=tl.float32)
 
 
 q_ptrs = q + pid_batch * (HEAD_DIM * N_HEAD) + pid_head * (HEAD_DIM) + tl.arange(0, HEAD_DIM)
 q_vec = tl.reshape(tl.load(q_ptrs), (HEAD_DIM, 1))

 context_len = tl.load(context_lens + pid_batch)
 max_block_index = tl.cdiv(context_len, BLOCK_SIZE) # 向上取整
 offs_kv = tl.arange(0, BLOCK_SIZE * HEAD_DIM)
 
 for block_idx in range(0, max_block_index):
 
 physical_idx = tl.load(block_table + pid_batch * MAX_BLOCKS_PER_SEQ + block_idx)
 physical_idx = tl.maximum(physical_idx, 0).to(tl.int64)
 base = (physical_idx * N_KV_HEAD * BLOCK_SIZE * HEAD_DIM + kv_head_idx* BLOCK_SIZE * HEAD_DIM)

 # 加载时 mask 掉超出 context_len 的 token
 token_start = block_idx * BLOCK_SIZE
 valid_in_block = tl.minimum(BLOCK_SIZE, context_len - token_start)
 kv_token_mask = tl.arange(0, BLOCK_SIZE * HEAD_DIM) < valid_in_block * HEAD_DIM

 k_ptrs = k_cache + base + offs_kv
 # k_block: (block_size, head_dim)
 k_block = tl.load(k_ptrs, mask = kv_token_mask, other=0.0)
 k_block = tl.reshape(k_block, (BLOCK_SIZE, HEAD_DIM))
 v_ptrs = v_cache + base + offs_kv
 v_block = tl.load(v_ptrs, mask = kv_token_mask, other=0.0)
 v_block = tl.reshape(v_block, (BLOCK_SIZE, HEAD_DIM))

 # mask 最后一个 block 的无效 token
 valid = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < context_len
 q_row = tl.reshape(q_vec, (1, HEAD_DIM)) # (1, HEAD_DIM)
 s = tl.sum(k_block * q_row, axis=1) # (BLOCK_SIZE,)
 s = s.to(tl.float32) * scale
 s = tl.where(valid, s, float('-inf'))

 m_new = tl.maximum(m, tl.max(s))
 alpha = tl.exp(m - m_new) # 旧的缩放因子
 p = tl.exp(s - m_new) # 当前 block 的权重

 l = l * alpha + tl.sum(p)
 p_col = tl.reshape(p, (BLOCK_SIZE, 1)) # (BLOCK_SIZE, 1)
 o = o * alpha + tl.sum(p_col * v_block, axis=0) # (HEAD_DIM,)
 m = m_new

 o = o / l # (1, HEAD_DIM)
 o_casted = tl.cast(o, out.dtype.element_ty)
 
 tl.store(out + pid_batch * (HEAD_DIM * N_HEAD) + pid_head * (HEAD_DIM) + tl.arange(0, HEAD_DIM), o_casted)

这部分本身没有什么概念上的困难。需要格外注意的是，调试的时候其报错是静默的：只会告诉你错在这里了，而不是告诉你这个Kernel里到底错在哪一行。因此，要多打印！

另一个非常要注意的关键点是Triton会默认执行死代码消除的优化（DCE）。这一点对于跑高性能是好事，但是对于调试极其不利。你可能会想，我注释掉后面一部分让Kernel报错的，一直二分直到找到问题所在，行不行呢？答案是不行，至少不是简单的行。如果你注释了后面一部分操作，而最后没有把前面操作的那部分store到某个地方的话，Triton编译器会直接把它们全部消除掉。那如果本身错误是发生在这部分，那么看上去你只是注释掉了最后几行就还是跑通了，但实际上并不能因此定位到错误所在。

踩坑点和注意事项

1、显存连续性问题

在第一版实现的时候，store_kvcache_kernel是没有传入stride参数的。这在Q、K、V三个权重矩阵分开计算的时候一点问题也没有。但是，在后来的计算流优化当中，为了减少Kernel的launch开销和反复读显存的开销，采用了一个很简单的优化实现：把$W_Q$、$W_K$、$W_V$拼接起来一起和输入$x$做GEMM。在两者同时进行的时候，出现了极其难以排查的乱码，几乎花费了三到五个小时才发现原因。这到底是为什么？看这段代码：

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


q, k, v = qkv.split([q_size, k_size, k_size], dim=-1)
q = q.view(B, seq_len, self.cfg.num_attention_heads, self.cfg.head_dim)
k = k.view(B, seq_len, self.cfg.num_key_value_heads, self.cfg.head_dim)
v = v.view(B, seq_len, self.cfg.num_key_value_heads, self.cfg.head_dim)

# RoPE
q, k = RoPE.apply_rope(q, k, cos, sin)
# print(v.is_contiguous())
# k/v reshape 成 (total_tokens, n_kv_heads, head_dim) 给 store kernel
k_flat = k.reshape(-1, self.cfg.num_key_value_heads, self.cfg.head_dim)
v_flat = v.reshape(-1, self.cfg.num_key_value_heads, self.cfg.head_dim)

# 写入 KV cache
store_kvcache(k_flat, v_flat, kv_cache_k, kv_cache_v, slot_mapping)# type:ignore

看上去没什么问题？

但如果把中间的print注释掉，你会发现，v.is_contiguous()的结果是false。如果改成k，结果又是true。这到底是为什么？因为：k经过了RoPE.apply_rope，在计算的过程中，pytorch自动帮你完成了向量的深拷贝，从而完成了向量的连续化重整，而v的两个操作，.view()和.reshape()都是不触发深拷贝的。也就是说，在它传入store_kvcache这个warpper的时候，它是仅仅是一个视图，在物理上不连续。

那这为什么会造成问题？因为之前说的：Triton内部的操作模式更接近C++：传入的Tensor实际上是隐式转换成带数据类型的指针来处理的。这意味着算术偏移不会考虑任何视图，操作的直接就是底层的物理地址。这和不连续明显是冲突的，但却又在整个大Tensor的数据范围内。因此就会发生，“数据不对，但也没有报错，看k的数据是完全对得上的，无论如何都查不出毛病”的问题。因此，pytorch的Tensor操作当中一定要小心视图和实际存储布局的区别问题。最后，是通过添加stride，重写底层的指针偏移来解决的。如果从头在C++本身开发，反倒没有这个问题了。

2、context_lens和seq_len的区别

这个我一时间还真不能完全精确的说明区别。为了避免对读者造成误会，直接请Gemini老师帮忙总结一下：

在手搓大模型推理引擎时，这两个变量在命名上极具迷惑性，如果不加区分，很容易在 Prefill 和 Decode 切换的瞬间触发经典的“差一错误”（Off-by-one），导致模型吐出的第一个新词直接变成乱码。简单来说，它们的物理含义完全不同：

seq_len (当前处理长度)：指的是模型当前这一步 Forward 到底吃进去了几个 token。
- 在 Prefill 阶段，一次性处理整个 Prompt，所以 seq_len = N。
- 在 Decode 阶段，因为我们是逐字生成，每次只喂给模型一个最新的词，所以 seq_len 永远等于 1。
context_lens (上下文总长度)：这是传给底层 Paged Attention 算子（如 Triton Kernel）的关键参数。它代表当前的 Query 到底需要和 多少个 token（包含它自己） 去计算注意力得分。

换言之，一个简单的自我思想检验：在Prefill刚刚结束，第一轮Decode开始的时候，传入Decode的参数应该是：seq_len=1，context_lens=N+1，其中N是prompt中的token数量。

3、复制粘贴的问题

孩子们，用ai帮忙写测试脚本的时候一定要看看自己复制了几遍，否则会发现自己的推理框架莫名其妙的总是会在Prefill结束之后多出一个token而找不到任何原因😭

性能Profiling

用nsys进行了GPU侧性能的Profiling，模式是eager，无编译的自动算子融合。结果如下：

Kernel 时间占比 实例数 单步耗时
──────────────────────────────────────────────────────────────
GEMM (gemvx, 主) 63.1% 1700 ~6.8ms
GEMM (gemvx, 次) 22.2% 560 ~2.4ms
cutlass WMMA (gate_up fused) 5.4% 560 ~0.6ms
RMSNorm elementwise 2.0% 3380 ~0.2ms
BinaryFunctor (residual add) 1.0% 1700 ~0.1ms
CUDAFunctor_add 0.9% 2240 ~0.1ms
Decode_Paged_GQAAttention 0.9% 560 ~0.1ms
RMSNorm reduce 0.9% 1140 ~0.1ms
CatArrayBatchedCopy 0.8% 1120 ~0.08ms
store_kvcache_kernel 0.2% 560 ~0.02ms
──────────────────────────────────────────────────────────────
GPU 总时间（20步） ~214ms 单步 ~10.7ms

按3.15GB（3GB权重+0.15GB的KV cache）的Memory Bound情况来计算，单步的理论最优化时间是：

$$ (3.15GB/token)/(382GB/s)=0.00825s=8.25ms $$

性能利用率已经达到77%，与此同时，整个单步Decoding的时延居然达到33ms左右，远远长于GPU侧的耗时。这主要是因为pytorch侧的时间非常长，反复调用产生了巨量的overhead。在我使用的WSL系统中，这个调用还会产生更大的虚拟化开销。因此目前的优化的性能瓶颈已经不是在GPU端了。这也是为什么接下来要做CUDA Graph。

Pico-vLLM 开发日志 #5 Paged Attention

Sat, 28 Mar 2026 01:22:15 +0800

终于到这一步了。刚开始开发Pico-vLLM这个推理框架的时候还觉得这是一件遥不可及的事情，但最终还是完成了。目前调用的forward过程里的Kernel仍然是伪Batching——内部是一个简陋的python for循环，不过整个Prefill、Decode、Engine、Scheduler、Forward的全流程都已经以完整的Paged Attention+Continuous Batching流程跑通了，只差Triton Kernel或者CUDA Kernel接入——我个人计划写Triton Kernel为主：毕竟时间成本有限，把性能提升到可以接受的程度是目前的目标。不论如何，这个里程碑还是挺有意义的。

和之前的博客中所提到的一样，Paged Attention就是专门为了解决内存碎片问题而生的。在naive kv cache的情况下，设计者必然面临两个决策中的一个：要么选择在一开始分配较短的kv cache空间，代价是随着seq_len长度超出原本预计的空间需要不断的进行完整的拷贝和重新分配；要么选择在一开始分配绝对够用的kv cache空间，代价是请求还没有产生几个，显存就迅速爆炸，而绝大部分大概率还没开始用。这一动机任何读者都应该相当了解，因此就不再冗余赘述了。

核心数据结构和设计模式

这次的实现规划的比较完整且一开始就预留了可扩展性，在最初就计划了双pool设计。也就是说，有可以防止硬OOM的GPU + CPU offload的双cache池，加上随之必须产生的逻辑block_id和物理block_id的分离。任何一个请求都会携带一个kv cache表，里面记录的是其持有的逻辑cache块的索引。而Block Manager则同时维护一个映射表：每个逻辑cache块对应的物理cache块的位置，以及物理块到底是哪种类型（CPU的还是GPU的还是未分配、非法的）的cache块。这种双重映射提供的分层抽象在可预见的未来甚至可以扩展到SSD——如果文件系统的性能允许的话。据作者本人所知，这种offload在工业界是大趋势，因此特意将其完整实现了出来。

核心代码块如下。值得注意的是，Block Manager本身不决定换入和换出的策略，仅仅是提供方法：这是思考后的结果，我个人认为这部分操作更适合分离开来，交给某个专门的Scheduler去执行，写在Block Manager内会让模式变得过度耦合且重量化。

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


class pagedblocktype(Enum):
 GPU = "gpu"
 CPU = "cpu"
 NONE = "none" # 表示未分配
class BlockManager:
 def __init__(self, 
 num_gpu_blocks: int, 
 num_cpu_blocks: int, 
 block_size: int, 
 num_layers: int,
 num_kv_heads: int, 
 head_dim: int, 
 dtype: dtype):
 ########## 物理块索引 ##########
 self.num_physical_gpu_blocks = num_gpu_blocks
 self.num_physical_cpu_blocks = num_cpu_blocks
 self.num_total_blocks = num_gpu_blocks + num_cpu_blocks
 self.block_size = block_size
 # 空闲块列表，初始全部空闲
 # 用 deque 比 list 的 pop(0) 快
 # GPU pool
 self.gpu_free_blocks: deque[int] = deque(range(num_gpu_blocks))
 self.gpu_kv_cache = torch.zeros(
 2, num_layers, num_gpu_blocks, block_size, num_kv_heads, head_dim,
 device='cuda', dtype=dtype
 )
 
 # CPU pool（offload 目标）
 self.cpu_free_blocks: deque[int] = deque(range(num_cpu_blocks))
 self.cpu_kv_cache = torch.zeros(
 2, num_layers, num_cpu_blocks, block_size, num_kv_heads, head_dim,
 device='cpu', dtype=dtype,
 pin_memory=True # ← 关键：pin_memory 让 CPU→GPU 传输更快
 )
 
 ########## 逻辑块索引 ##########
 # 逻辑块表，记录每个逻辑块对应的物理块和类型
 # 最多逻辑块不会超过 num_total_blocks，因为每个逻辑块至少占一个物理块
 # 第i个初始逻辑块初始化为[NONE, -1]，表示未分配物理块
 self.block_mapping: List[tuple[pagedblocktype, int]] = [
 (pagedblocktype.NONE, -1) for i in range(self.num_total_blocks)
 ]
 self.logical_free_blocks: deque[int] = deque(range(self.num_total_blocks)) # 逻辑块索引，初始全部空闲
 
 def allocate(self, num_blocks: int = 1) -> List[int]:
 # 分配 num_blocks 个物理块
 block_ids = []
 if num_blocks > self.num_free_blocks:
 raise RuntimeError(f"Not enough free blocks to allocate {num_blocks} blocks")
 for _ in range(num_blocks):
 if self.gpu_free_blocks:
 physical_block_id = self.gpu_free_blocks.popleft()
 logical_block_id = self.logical_free_blocks.popleft()
 self.block_mapping[logical_block_id] = (pagedblocktype.GPU, physical_block_id)
 block_ids.append(logical_block_id)
 elif self.cpu_free_blocks:
 physical_block_id = self.cpu_free_blocks.popleft()
 logical_block_id = self.logical_free_blocks.popleft()
 self.block_mapping[logical_block_id] = (pagedblocktype.CPU, physical_block_id)
 block_ids.append(logical_block_id)
 else:
 raise RuntimeError("No free blocks available")

 return block_ids

 def free(self, block_ids: List[int]) -> None:
 # 回收物理块，加回 free_blocks
 # 回收逻辑块
 for block_id in block_ids:
 block_type, physical_block_id = self.block_mapping[block_id]
 if block_type == pagedblocktype.GPU:
 self.gpu_free_blocks.append(physical_block_id)
 elif block_type == pagedblocktype.CPU:
 self.cpu_free_blocks.append(physical_block_id)
 else:
 raise RuntimeError(f"Block {block_id} is not allocated")
 # 更新 block_mapping
 self.logical_free_blocks.append(block_id)
 self.block_mapping[block_id] = (pagedblocktype.NONE, -1)

 #####################################################
 # 非必须的方法，用于offload到CPU，swap_in/swap_out()
 #####################################################
 def swap_out(self, block_ids: List[int]) -> None:
 """
 把 GPU 上的 block 换出到 CPU
 """
 cpu_block_ids = []
 for block_id in block_ids:
 block_type, physical_block_id = self.block_mapping[block_id]
 if block_type != pagedblocktype.GPU:
 raise RuntimeError(f"Block {block_id} is not on GPU, cannot swap out")
 cpu_block_id = self.cpu_free_blocks.popleft()
 self.cpu_kv_cache[:,:,cpu_block_id].copy_(
 self.gpu_kv_cache[:,:,physical_block_id], non_blocking=True
 )
 self.gpu_free_blocks.append(physical_block_id) # GPU 块释放
 self.block_mapping[block_id] = (pagedblocktype.CPU, cpu_block_id) # 更新映射
 cpu_block_ids.append(cpu_block_id)

 def swap_in(self, block_ids: List[int]) -> None:
 """
 把 CPU 上的 block 换回 GPU
 """
 gpu_block_ids = []
 for block_id in block_ids:
 block_type, physical_block_id = self.block_mapping[block_id]
 if block_type != pagedblocktype.CPU:
 raise RuntimeError(f"Block {block_id} is not on CPU, cannot swap in")
 gpu_block_id = self.gpu_free_blocks.popleft()
 self.gpu_kv_cache[:,:,gpu_block_id].copy_(
 self.cpu_kv_cache[:,:,physical_block_id], non_blocking=True
 )
 self.cpu_free_blocks.append(physical_block_id) # CPU 块释放
 self.block_mapping[block_id] = (pagedblocktype.GPU, gpu_block_id) # 更新映射
 gpu_block_ids.append(gpu_block_id)

 @property
 def num_free_blocks(self) -> int:
 return len(self.gpu_free_blocks) + len(self.cpu_free_blocks)

 def can_allocate_gpu(self, num_blocks: int) -> bool:
 # 查询是否有足够的GPU空闲块
 # scheduler 用这个决定是否接受新请求
 return len(self.gpu_free_blocks) >= num_blocks
 

而从这个Block Manager里取回相应的物理块的index组成的block_table的方法则如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


 def get_block_table(self) -> torch.Tensor:
 """
 读出当前所有已写入的 k/v
 return: 存储的所有对应index的物理index
 """
 physical_ids = []
 for lid in self.cache_block_index[:self.allocated_cache_block_num]:
 _, pid = self.block_manager.block_mapping[lid]
 physical_ids.append(pid)
 return torch.tensor(physical_ids, dtype=torch.int32, device=self.device)

现在它是变长的，但随后的Triton Kernel当中，它可能会变成定长的——取决于到时候的性能调试了。

此外特别需要注意的是自回归过程中，关于真正的“序列长度”的相关处理。虽然seq_len不包括自己正在生成的这一个token——它的Logits还不存在呢，但是在每一层layer当中，其实k和v已经完整存在了，是按照这个token对应的k、v存在去处理的。因此在做相关操作的时候要格外小心，在适当的地方进行+1的修正，并且提前做好kv cache块的预申请和预分配，以免发生问题。

再此外，传参时候的类型提示真的很重要。虽然python不是静态类型的强类型语言，但适当的类型提示对于编程的帮助极其极其大——最大的之一就是你可以随时随地的查看这个变量是什么，以及通过自动补全查看和迅速编写相应的方法，避免无穷无尽的来回翻找。嗯，果然还是Java/C/C++好啊。

吐槽

作者本人用的卡是5070，著名的nano-vllm项目在写的过程中偷奸耍滑，直接用了flash-attn库——作者本人原本计划同样用flash-attn的算子接入跑通全流程之后作为对比的benchmark，测试自己的算子的信念，但尝试在wsl的venv环境中以各种方法安装flash-attn库均失败。通过github和ai查询原因，得到的结论是，5070是Blackwell系列的消费卡，而目前flash-attn目前似乎是只支持hopper以及之前的架构——直接没法用了（至少截止2026年3月28日不行。此结论有时效性，读者请注意）。这进一步倒逼作者去真的写一个Triton Kernel，不过总归是要写的，最后都一样嘛。

大家如果想使用这些算子库，切记查看好仓库发布页的目前支持什么cuda版本、pytorch版本、显卡系列的提示，避免做无用功。

性能数据

令人意外的是实现了如此多冗余拷贝的naive paged Attention之后性能居然不降反升（3%），这似乎是得益于forward过程没了cat()，某种程度上反而同时减少了拷贝。虽然之前说baseline是不适用kv cache的彻底naive实现，但那太离谱了，这次的才算真正的有比较意义的baseline。在接下来一段时间里，就主要尝试对这个指标做提升了，尝试在有限的时间内逐步优化到vLLM的60%水平，然后再迈进到多机阶段。

具体数据直接看图⬇️

Pico-vLLM 开发日志 #4 Continous Batching

Thu, 26 Mar 2026 20:41:19 +0800

今天的主要成果是写了一个pseudo的Continous Batching框架。pseudo的意思是，engine内部的调用其实仍然是按照loop的模式一个个进行，但scheduler的功能本身是完善的，step模式的完整流程也已经写完。这么做主要是因为paged Attention本身作为框架的必须组件还没有实现——在实现paged Attention的过程中，从GQAAttention算子的模式本身到engine到KV cache的cache类和BlockManager全部都得重构或者从头开始写——因此在一个必然要推翻的框架内写一版必然要在接下来几天内全部重构的代码没什么意义，不如接下来直接实现完整的既有paged又有batching的版本。

毕竟我又没有交付压力🤓。

Continous Batching是什么

是一种可以增加推理框架调度效率的方法。在这个技术普及之前，Batching作为提高推理效率的技术已经被广泛应用。因此，要首先理解Batching在推理当中为什么可以提高推理效率。

一些博客或者技术分享可能会说Batching能提高算术强度什么的，但这其实是一种误解，提高算术强度的技术是speculative decoding，而不是Batching技术。这是因为Batching产生的矩阵其实是对角线分块的：互相并不产生计算交叠，反而会产生大量0的空洞，这一点（在naive Kernel里）通过Attention mask遮盖。同时，KV cache本身的访问也是不会减少的，每一个之前算过的依然要完整读一遍。它最大的意义其实在于减少了权重的读取次数——原本计算一个sequence，吐一个token就要从主存（低层次的存储层次）读一遍完整的权重到高层次当中，现在计算$N$个sequence才需要读一遍，中间的过程均驻留在高层次存储里不离开。对于权重很大的大模型来说，这个效率提升其实是非常重要的。从这个角度来说，Continous Batching计的优化思想根基和GEMM矩阵乘法的优化也没有根本性的差别——计算机基本原理，加速大概率事件和有效利用时空局部性的又一个良好体现。

Batching产生了另一个问题：不同的sequence不太可能是同等长度的，而且不太可能同时吐出eos，到达停止状态。将类似长度的sequence打包在一起可以部分地解决第一个问题，但不能解决第二个问题，这个问题同样被后面要做的paged Attention大部分地解决。这就导致一个结果：在一个batch中的一部分sequence已经吐出了eos，变成停止状态了，而另外一些则没有。只要一个sequence还在跑，整个batch是不能停止的。这个时候，batch可能反而变成负收益：许多读写已经是完全空转的了，只是为了陪跑还在自回归过程的sequence，但内存读写消耗反而一点也没有减少。Continous Batching解决的是第二个问题。

所以实际上，Continous Batching和paged Attention本就是深度耦合的两个技术：如果没有paged Attention，Continous Batching每次都要完整的重新分配和整理内存，那么Continous Batching本身就毫无意义了。反过来说，没有Batching和Continous Batching，在单序列请求下，内存碎片这个问题可能还没有那么严重，paged Attention的威力也完全无法有效发挥。试想：在有了paged Attention之后，Continous Batching其实完全不需要进行任何内存操作：只需要把换出的sequence的KV cache页表索引换掉，变成下一个sequence的就可以了，也就是说对sequence本身其实无感知。这也难怪这两件事被同一时间提出，作者在研究的过程中很可能遇到了这个问题然后把它想明白了。

Continous Batching的设计、实现和调度策略

其实这个技术的基本实现也在思路上比较顺畅。它的基本思想是把所有sequence分成waiting、prefilling、decoding和finished四种状态（也有的实现是waiting、running和finished三种，取决于你如何写你的推理engine和kernel，两者是耦合的），以及engine的step模式设计。通过一个scheduler的调度，scheduler管理所有请求、赋予每一个请求一个unique的id编号，并且决定下一个step处理哪些，如何改变序列的状态，并且维持每个step内的负载大致均衡。engine则负责进行prefilling、decoding，并且检查decoding序列中的情况，将已经达到max_len的和已经生成出eos的序列中的sequence移出，进入finished，而finished则负责记录生成的结果、等待外部请求获取，并且提供将finished队列清空的方法，实际上是为了方便外部请求的设计。

在这个过程中有意思的反而是scheduler的调度策略设计：由于prefilling和decoding的计算资源需求并不相同（其实这里已经初步开始产生PD分离的思想了：异质化的负载调度总是比较难处理的），scheduler如何同时确保效率和公平性就是一个问题。prefilling一次应该多少？decoding一次应该多少？这个问题相当open，理论上来说应该是相当值得研究的（尤其是对于生产环境、ToC环境来说，因为它会极大的影响TTFT指标）。在目前的版本中，为了简单起见，实现的是先到先服务（FCFS）和Prefilling+Decoding中的sequence数量总和固定。这是一个粗糙的调度，不过对于框架的可运行性没有影响。也许后面可以考虑把策略拿出来单独设计。

Continous Batching的相关代码

主要的类有两个，分别是Request和Scheduler。Request记录一个请求的基本信息，Scheduler基于前者的基本信息进行调度。需要注意的是，在我的框架里我本人选择了把KVcache和Request放在一起管理。其他框架似乎不是这样，但按这个模式似乎也没什么问题。

Request类：

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


class RequestStatus(Enum):
 WAITING = "waiting" # 在 waiting 队列，还没做 prefill
 PREFILL = "prefill" # 本步正在做 prefill
 DECODING = "decoding" # 已经 prefill 完，正在 decode
 FINISHED = "finished" # 已完成

''' 请求对象，包含请求的所有信息和状态
- request_id: 请求 ID，唯一标识一个请求
- input_ids: 输入的 token ids，shape (1, init_seq_len)，包含整个 prompt
- generated_ids: 已经生成的 token ids，shape (1, )，初始为 空
- max_new_tokens: 最多生成多少个 token
- temperature: 采样温度，传递给 sampler
- top_p: top-p 截断，传递给 sampler
- kv_cache: 每个请求独享一个 KV cache 实例，存储生成过程中的 KV 状态
'''
class Request:
 request_id: int
 input_ids: List[int] # (1, init_seq_len)，包含整个 prompt
 generated_ids: List[int] # (1, )，包含已经生成的 token ids，初始为 空
 max_new_tokens: int
 temperature: float
 top_p: float
 kv_cache: Optional[KVCache] # 每个请求独享一个 KV cache 实例
 request_status: RequestStatus
 has_eos_token: bool # 是否已经生成 eos_token，scheduler 不直接接触 tokenizer 和 eos_token_id，这个由 engine 在 decode_step 后更新

 def __init__(self, request_id: int, input_ids: List[int], max_new_tokens: int, temperature: float, top_p: float):
 self.request_id = request_id
 self.input_ids = input_ids
 self.generated_ids = [] # 初始为 空
 self.request_status = RequestStatus.WAITING
 self.max_new_tokens = max_new_tokens
 self.temperature = temperature
 self.top_p = top_p
 self.kv_cache = None
 self.has_finished_notification = False # engine改变这个状态，scheduler根据这个状态改变 request_status和移出队列

 def is_max_len_finished(self) -> bool:
 return len(self.generated_ids) >= self.max_new_tokens

 @property
 def prompt_len(self) -> int:
 return len(self.input_ids)
 
 @property 
 def total_len(self) -> int:
 return len(self.input_ids) + len(self.generated_ids)

Scheduler类：

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


class Scheduler:
 _next_request_id: int # 用于生成唯一的 request_id
 max_num_seqs: int # 同时处理的最大请求数，超过这个数的新请求会排队等待
 kv_cache_cls : type[KVCache] # KVCache 的类，用于创建请求的 KV cache 实例
 kv_cache_kwargs = {} # 创建 KV cache 实例的参数字典

 waiting: List[Request] # 还没开始处理的请求
 prefilling: List[Request] # 正在 prefill 的请求
 decoding: List[Request] # 正在 decode 的请求
 finished: List[Request] # 已完成的请求
 def __init__(self, kv_cache_cls: type[KVCache], kv_cache_kwargs={}, _next_request_id: int = 0, max_num_seqs: int = 4):
 self._next_request_id = _next_request_id
 self.max_num_seqs = max_num_seqs
 self.kv_cache_cls = kv_cache_cls
 self.kv_cache_kwargs = kv_cache_kwargs
 self.waiting = []
 self.prefilling = []
 self.decoding = []
 self.finished = []

 ''' 插入新请求
 - input_ids: 输入的 token ids，shape (1, seq_len)，包含整个 prompt
 - max_new_tokens: 最多生成多少个 token
 - temperature: 采样温度，传递给 sampler
 - top_p: top-p 截断，传递给 sampler
 - cache_cls: KVCache 的类，用于创建请求的 KV cache 实例
 - cache_kwargs: 创建 KV cache 实例的参数字典
 return: request_id，唯一标识一个请求'''
 def insert_request(self, 
 input_ids: List[int], 
 max_new_tokens: int, 
 temperature: float, 
 top_p: float, ) -> int:
 request = self.create_request(input_ids, max_new_tokens, temperature, top_p, self.kv_cache_cls, self.kv_cache_kwargs)
 self.add_request(request)
 return request.request_id

 '''' 创建新的请求对象
 - input_ids: 输入的 token ids，shape (1, seq_len)，包含整个 prompt
 - max_new_tokens: 最多生成多少个 token
 - temperature: 采样温度，传递给 sampler
 - top_p: top-p 截断，传递给 sampler
 - cache_cls: KVCache 的类，用于创建请求的 KV cache 实例
 - cache_kwargs: 创建 KV cache 实例的参数字典
 return: request_id，唯一标识一个请求'''
 def create_request(self, 
 input_ids: List[int], 
 max_new_tokens: int, 
 temperature: float, 
 top_p: float, 
 cache_cls, 
 cache_kwargs) -> Request:
 request_id = self._next_request_id
 self._next_request_id += 1
 return Request(request_id, input_ids, max_new_tokens, temperature, top_p)

 '''' 添加新请求到等待队列
 - request: 新的请求对象
 '''
 def add_request(self, request: Request):
 self.waiting.append(request)

 ''' 调度器主循环
 return: 2个列表，分别是当前处于 prefilling、decoding 状态的请求列表
 - prefilling: 正在 prefill 的请求
 - decoding: 正在 decode 的请求
 '''
 def schedule(self) -> tuple[List[Request], List[Request]]:
 
 # 1. 对 decoding 队列中的请求进行 decode，完成后移动到 finished 队列
 for request in self.decoding[:]:
 # is_max_len_finished 由 scheduler判断
 # has_finished_notification 由 engine 在 decode_step 后更新，解耦两者的逻辑，scheduler 不直接接触 tokenizer 和 eos_token_id
 if request.has_finished_notification == True:
 request.request_status = RequestStatus.FINISHED
 self.decoding.remove(request)
 self.finished.append(request)

 # 2. 对 prefilling 队列中的请求进行 prefill，完成后移动到 decoding 队列
 for request in self.prefilling[:]:
 # prefill 完成后：
 request.request_status = RequestStatus.DECODING
 self.prefilling.remove(request)
 self.decoding.append(request)
 
 # 3. 从 waiting 队列中取出请求，放入 prefilling 队列，直到达到 max_batch_size
 # 不进行kv_cache的创建和初始化
 while self.waiting and self.num_in_progress < self.max_num_seqs:
 # 先到先服务策略
 request = self.waiting.pop(0)
 request.request_status = RequestStatus.PREFILL
 self.prefilling.append(request)
 # kv cache必须定长，否则报错
 # request.kv_cache = self.kv_cache_cls(**self.kv_cache_kwargs)

 return self.prefilling, self.decoding

 '''' 清空 finished 队列，释放资源'''
 def clear_finished(self):
 self.finished.clear()

 def get_running_requests(self) -> List[Request]:
 return self.prefilling + self.decoding
 
 @property
 def num_waiting(self) -> int:
 return len(self.waiting)
 @property
 def num_prefilling(self) -> int:
 return len(self.prefilling)
 @property
 def num_decoding(self) -> int:
 return len(self.decoding)
 @property
 def num_in_progress(self) -> int:
 return len(self.prefilling) + len(self.decoding)
 @property
 def num_finished(self) -> int:
 return len(self.finished)

值得指出的是id用int是有明确缺陷的：str和uuid作为id很可能是更好的选项。在这里用int只是因为它简单，而且我的卡大概也不会有超过2147483647个请求。

测试一遍通过，这里就不贴了。接下来就是实现真正的paged Attention了，个人认为它的意义可能不太亚于KV cache本身，直接开启了一小个时代级别的想法。

Pico-vLLM 开发日志 #3 KV cache

Wed, 25 Mar 2026 01:15:40 +0800

今天实现了naive的KV cache，其效果之好确实令人震惊。vLLM根本没有提供不使用KV cache的推理方案，而且会按比例占据所有的显存，足可见KV cache在推理框架中的核心地位。先前的实验也证明了，如果不使用KV cache，推理速度很快就会下降到一个令人难以忍受的程度，从而不允许进行任何有实际意义的模型推理。

KV cache的思想

KV cahce的思想很简单，就是空间换时间。通过保存每一个先前token在经过各层（以qwen2.5-1.5B为例，就是28层）的attention层的时候，把它们的副产物：K和V保存下来。所以实际上，KVcache的shape大约是这样的（对一个请求来说）：

$$ (Layer\_Index, K(0)/V(1), Seq\_Len, Num\_KV\_Heads, Head\_dim) $$

在自回归过程中，一次前向传播，每经过一个TransformerBlock，前向传播的过程就把计算出的结果保存到一个大KV cache的相应位置上。Prefill的过程一次保存等同于input_prompt的token数量的KVcache，decoding过程则一次完整的前向传播增加一个。

虽然思想很简单，但工程实现相对来说没有那么简单。因为要对整个engine进行重构：原本不需要区分Prefill和Decoding，但现在需要了。进一步的说，这两者的需求根本不同：对于Prefill，由于方阵形状规整，算术强度是比较高的，因此Roofline模型更靠近compute bound那一端；而对于decoding来说几乎所有的时间都用来了进行模型权重的读取和KV cache的读取（虽然在我们这个例子当中效果并不显著），因此完全是memory bound的。这样，确实是从实践的角度感受了一下PD分离为什么最终会发生，这两者的计算模式是非常不一样的，尤其是当初始input的prompt很长的时候。

除此之外，一个非常容易错的点是：第一次写的时候很容易在每一层TransformerBlock里都把KVcache的Seq_len加一，但这样会造成同一个token的不同层的KVcache被存在了完全错位的地方。实际上，在layer_idx = max_layer - 1的时候把Seq_len加一就可以了。这是一个非常容易实现不对的地方，当时debug了半天。最后的代码长这样：

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


# 朴素实现
class NaiveKVCache(KVCache):
 # 连续 tensor
 def __init__(self, num_layers, max_seq_len, num_kv_heads, head_dim, device, dtype):
 self.num_layers = num_layers
 self.max_seq_len = max_seq_len
 self.num_heads = num_kv_heads
 self.head_dim = head_dim
 self.device = device
 self.cache = torch.zeros(
 num_layers, 2, max_seq_len, num_kv_heads, head_dim,
 device=device, dtype=dtype
 )
 self._seq_len = 0 # 当前已填入的长度

 def update(self, layer_idx: int, k: Tensor, v: Tensor) -> None:
 # k, v: (1, num_kv_heads, head_dim) 或 (seq_len, num_kv_heads, head_dim)
 # 通过读取kv的shape自动适配 prefill 和 decode 场景
 # 以后用 cuda kernel 的时候这个要怎么处理？可能需要分开 prefill 和 decode 的接口，或者额外传参？
 num_new_tokens = k.shape[0]
 start = self._seq_len
 end = self._seq_len + num_new_tokens

 self.cache[layer_idx, 0, start:end, :, :] = k
 self.cache[layer_idx, 1, start:end, :, :] = v
 if (layer_idx == self.num_layers - 1): # 只在最后一层更新 seq_len，保证所有层的 seq_len 一致
 self._seq_len += num_new_tokens

 def get(self, layer_idx: int) -> tuple[Tensor, Tensor]:
 return self.cache[layer_idx, 0, :self._seq_len, :, :], self.cache[layer_idx, 1, :self._seq_len, :, :]

 def reset(self) -> None:
 self._seq_len = 0

 @property
 def seq_len(self) -> int:
 return self._seq_len

希望看到的人能够绕过这个坑。

下面是使用了naive实现的KV cache之后，模型的decoding性能随着seq_len的变化的profiling图。

可以看到，对于1.5B这样的小模型（权重3GB）和<4096的seq_length长度来说，对于不同的sequence长度，单个token的推理时间几乎是一个常数，不会发生明显下降，和之前的naive实现的$O(N^2)$的曲线特征形成了非常明显的对比。这主要是因为权重读取相关的操作占据了绝大部分：这部分操作和seq_len是完全无关的，是$O(1)$的常数项。如果把length拉长，应该会看到后续的$O(N)$阶段，主要由KV cache本身的读取导致的带宽需求决定的Latency。不过由于我的卡（5070）的显存很小（8GB），在整个显存能够不OOM的范围内，几乎都不会发生这种情况就是了。

理论上来说，5070能够达到的推理速度大约是96tokens/s，目前的利用率仅仅只有22%左右。可以发现能够优化的空间还是相当大的，Kernel Fusion能够有相当的发挥空间。

一些思考

另外，今天和同学交流了一下未来的专精方向。算子开发本身似乎已经比较成熟了，通信库、高性能网络更值得做一些。确实如此，triton等等已经大大降低了开发难度，而且flashattn等各种算子融合和实现的库在pytorch里的生态也非常完善了。然而我的手上没有相关资源，只有单卡。但不论如何，也许应该考虑调整框架的开发方向？

目前的规划依然是要实现完整的KV cache。然而，prefix cache在目前的条件下，似乎没有个人开发者环境下有效评测的方法，这部分应该没有特别的高优先级的需求——可以往后推一推。

这样的话，目前的规划就是这样的：1、实现paged attention，2、进行Kernel Fusion的尝试和优化，3、实现单节点多卡推理，4、实现多节点多卡推理。

Pico-vLLM 开发日志 #2 自回归和Profiling

Tue, 24 Mar 2026 00:12:41 +0800

今天凌晨完成了自回归框架的初步搭建，到目前为止可以完整的进行模型的自回归生成了！如下图：

值得纪念的时刻。需要注意到的是：Greedy的sample策略会很容易引发一个大家现在已经不太关注的问题：注意力塌陷，即模型开始反复输出同一段话，无法进行任何有效的长文本输出。如果使用带有随机性的采样，这种情况就基本上不会发生了。这应该也是为什么Greedy在现在的模型采样方法里完全不采用。

而且可以大概看出，在短Prompt下，一个token生成的速度大约是~20tokens/sec。可以计算一下理论上其应该有的计算峰值：

$$ 理论最大速度 = 带宽 / 权重大小 = 288 GB/s / 3 GB ≈ 96 tokens/sec $$

可以看出，naive实现相对于理想情况来说有效利用的比例并不高，只达到了大约21%的理论峰值。接下来我们的目标就是提升这个数字，并且提升OOM的阈限（利用主存甚至SSD）——先画个饼，提升到60%以上？

下面进行profiling，这在后面的开发当中将会作为baseline反复用到，用于和后面每一次优化进行比较。从16个token开始，每次增加16个token长度的seq_len，一直增加到4096。profiling的结果如下所示：

模式：Naive（无 KV Cache）硬件：RTX 5070 Laptop 8GB

Seq_len	ms/step	token/s
16	47ms	21.0
512	55ms	18.2
1024	115ms	8.7
2048	272ms	3.7
4096	764ms	1.3

可以明显的看到，在纯naive实现下，其主要分两段：常数段（生成速度几乎不变）和二次函数段（生成速度减慢的模式完美符合二次函数的增长）。常数段我认为应该主要是因为Latency以模型权重的反复加载为主，compute的时间几乎可以忽略不计，因此前向传播的Latency是固定的。而二次函数段就是经典的$O(N^2)$的注意力计算复杂度compute-bound导致的$O(N^2)$时间增长。此外特别值得注意的是，在大约seq_len=512~528的地方，出现了非常明显的性能阶跃式的下降。目前考虑的有pytorch的kernel调度的不同导致的，也有可能是L2cache溢出导致的换页导致的。究竟是什么导致的，后续也需要仔细研究。

总的来说，如果不使用KV cache的话，内存其实远远达不到OOM的程度（8+16GB的显存，仅仅用了4GB），token的生成速度就已经慢到无法接受了，更别提什么CUDA kernel的优化，都是于事无补。

此外，在测试的过程当中，第一次profiling的implementation在理论显存占用仅仅3.3GB不到的时候就OOM了。然而，反复计算理论显存也没找到为什么OOM。查询了之后才知道，原来pytorch不会自动释放显存，会把已经不再用到的、原先用过的显存块作为reserved模式保留在内部，哪怕马上就要OOM了也不会自动释放（确实这么离谱）。而每次因为序列是越来越长的，原先reserved的内存从未用上过，造成碎片单向增加而从不减少。这某种意义上来说和paged attention遇到的和解决的问题非常相似：只不过pytorch这个管理的更加低级。后面将会逐步实现naive的KV cache和Paged attention机制，随后再考虑kernel的优化。

Pico-vLLM 开发日志 #1 Sampler的设计

Sun, 22 Mar 2026 23:37:49 +0800

今天的主要任务是写sampler，分为以下几种：Greedy，Temperature，Top-P。其实理论上来说Top-p和Temperature两个方法是解耦的，实现里的Top-p其实是Top-p+Temperature。不过无论哪种都很好实现，只讲基本思想。

Greedy，顾名思义就是永远只选择概率（Logits）最大的那个token。数学公式如下：

$$ X_{next}=argmax_i(Logits(X_i)) $$

Temperature是针对Greedy的改进，想法也很简单，就是用一个参数决定分布的平均程度（热化程度），然后用这个分布采样。由于实际上是模拟玻尔兹曼分布，和热力学关系很大，索性就叫Temperature了。很明显，T越大采样越随机，=1则和原本的随机采样分布相同。它的数学公式如下：

$$ \tilde{L}_i = \frac{\text{Logits}(X_i)}{T} $$

$$ P(X_i) = \text{Softmax}(\tilde{L}_i) = \frac{e^{\tilde{L}_i}}{\sum_j e^{\tilde{L}_j}} $$$$ X_{next} \sim P(X_i) $$

Top-P可以不包含T，但一般包含T。它是对Top-K（取固定前K个）的改进：取若干个最有可能的，直到覆盖了原本分布的前p部分，随后在这部分里面采样，以规避过于离谱的输出的同时又保证随机性。在包含T的情况下，它的公式如下：

$$ \mathcal{V} = argmin_{S \subseteq \mathcal{V}_{vocab}} |S| \quad \text{s.t.} \sum_{X_i \in S} P(X_i) \geq p $$

在候选集内重新归一化后采样：

$$ P'(X_i) = \begin{cases} \dfrac{P(X_i)}{\sum_{X_j \in \mathcal{V}} P(X_j)} & X_i \in \mathcal{V} \\ 0 & \text{otherwise} \end{cases} $$

$$ X_{next} \sim P'(X_i) $$

原理应该说很简单，但top-p算法在直觉上似乎有很大的优化空间，$O(NlogN)$的复杂度对于大Vocabulary的模型的推理来说会不会产生比较明显的性能影响？这需要全部跑通之后进行profiling，现在先不管了。

Pico-vLLM 开发日志 #0 博客安装

Sat, 21 Mar 2026 00:00:00 +0000

HelloWorld!

今天的任务是先把博客搞好，然后把刚刚跑通能吐一个token的框架py代码推送。

测试一下

测试成功了，后面就在这里更新。目前的计划是先跑通自回归的引擎，然后做最基本的profiling，随后逐步优化。具体的实现可以参考vllm和sglang，另外预计也会有很多写py的过程中想到的可以优化的点，后面都会整理在这里。

VLLM on Fain的Blog

Pico-vLLM 开发日志 #14 Prefix Caching

架构设计：在 vLLM 的地基上建 SGLang 的房子

vLLM模式

SGLang模式

粒度不一致的桥接

三层架构

insert 的 block 归属

语义设计：处理边界情况和分层解耦

两层引用计数的必要性

lock_ref 的 node 指针方案

Bug 1：release 用 match 重新找 block

Bug 2：split 后边长不对齐 block_size

Bug 3：token 路径遍历的根本缺陷

修复方案：node指针方案

权责清单

设计模式总结，SGLang vs vLLM vs Pico-vLLM

Benchmark和性能测试

短 Prompt（144 token）

长 Prompt（2083 token shared + ~35 token variable）

总结和心得

Pico-vLLM 开发日志 #13 PD分离（续） 同步异步模式

什么是同步和异步模式

同步模式

异步模式

NIXL的必要性

生产者-消费者模型

基于异构的PD分离-TP并行设计

KV Cache的重分片

Pico-vLLM 开发日志 #12 PD分离

为什么需要PD分离

TTFT和ITL的权衡

硬件资源的充分利用

可用的工具箱

使用torch.distribute（NCCL）

使用NIXL

使用CUDA IPC

设计模式

基于实例角色role的兼容式设计

传输层抽象分离

注意事项和Bug

值得注意的Bug汇总

1、Python的可变默认值问题

2、锁死问题

Pico-vLLM 开发日志 #11 张量并行（续） 数据对比分析

性能数据汇总

wall-clock端到端时间（含CPU开销）

纯GPU replay时间（无CPU开销）

AllReduce 延迟（单次3KB传输）

值得一提的细节

Pico-vLLM 开发日志 #10 张量并行 TensorParallel

张量并行涉及到的改动

模型

KV cache的存储

引擎Engine

张量并行的实现手段

需要注意的技术细节

张量并行在Qwen2.5-1.5B上的效果

画饼

LLM 学习日志 #2 集合通信拓扑、原语和实现

集合通信的原语

点对点（P2P）

Broadcast

Reduce

All-Reduce

Reduce-Scatter

All-Gather

All-to-All

All-to-All-v

集合通信的常见拓扑

Ring（环形）

Tree（树形）

Double Binary Tree

Recursive Halving-Doubling

集合通信的常见实现

Cost Model：运行时选择

硬件互联的数感建立

一、GPU 节点内互联：NVLink

二、GPU 节点内互联：NVSwitch

三、CPU-GPU 互联：PCIe

Pico-vLLM 开发日志 #13 PD分离（续）同步异步模式

Pico-vLLM 开发日志 #11 张量并行（续）数据对比分析