博士延期一年是坏事吗？不，是好事。
做开源项目吃力不讨好？不，这是难得的机遇。
最近看到Kaichao在知乎上分享了他与vLLM团队的经历，我也忍不住想写点什么。这不仅是对过去一年的总结，也是对我博士生涯最后一段时光的记录。2024年，我选择多读一年PhD，投身大语言模型的研究，并与vLLM团队一起，开启了一段充满挑战与机遇的旅程。
为什么选择vLLM？
2023年，我还在研究如何压缩视频，确保神经网络能准确统计车辆数量。但随着时间推移，这个领域最具创新性的想法已被前人探索殆尽，剩下的工作大多是在不同视频中调整压缩参数，感觉像是在螺蛳壳里做道场。与此同时，ChatGPT的横空出世，让大语言模型（LLM）系统的研究迎来了爆发式增长。我个人平时也喜欢生成文字，这和大语言模型听起来就很配。于是，本可以准备毕业的我决定多读一年PhD，转向大语言模型研究，并请导师帮我联系暑研。
与vLLM的缘分
我的导师作为他导师的关门弟子，帮我联系上了他导师的开山弟子 Ion Stoica（在此特别感谢我导和 Ion）。巧合的是，Ion 正是 vLLM 团队的指导教授。就这样，我与 vLLM 团队的故事开始了。
第一关：面试
在正式加入vLLM团队之前，Ion简短的面试了我。面试开始我感受到了久违的压力：一开始会议链接点不进去，我紧急排查是不是链接出错了（事后发现是Google meet的问题），然后赶紧和Ion发邮件改用zoom，然后反复刷新邮件等Ion的回复。好在最后面试总算成功开始。
其中一个让我印象深刻的问题是：描述你最骄傲的一个项目。在脑子空白之际，我在自己最有影响力的项目和最喜欢的项目（同时基本没有影响力）之间闪烁，最终凭借直觉选择讲我最喜欢的项目。在讲这个项目的过程中我文思泉涌，整个人都活了过来，最终通过了面试。
面试之后，我琢磨这个面试问题，突然发现了一种遥远的相似性：Ranveer（MSR一个组的老大）说做出最顶尖的研究需要的是和研究项目之间情感上的链接；我在讲述自己喜欢的项目时，眼睛都会放光；Haijie（Conviva程序员们的头头）在挑选人才的时候，最想看到的是motivation。
如果用一个词来概括这种情感连接/眼中的光芒/motivation，我觉得不是“动机”，而是“热情”。做出最好的项目，需要真正的热情。这种热情只能从生活中获得，写不进简历，但有了热情，人就有了精气神。
初入vLLM：性能基准测试
2月完成面试后，我以最快速度做好准备，4月顺利抵达伯克利。
不记得听谁说起过，系统研究的第一步都是做性能基准测试(performance benchmark)。所以到了vLLM团队之后，我做的第一件事就是建立一系列的性能基准测试。
在和vLLM团队的大家聊天过后，我发现做性能基准测试要展示给两类人看，背后是两种截然不同的动机和要求。
- vLLM的开发者们想看性能基准测试，目的是为了了解vLLM的性能到底有没有因为新贡献的代码变好或者变坏。这类benchmark在每次代码变动时都要跑一次，所以不能跑太久。
- LLM从业者们想看性能基准测试，目的是为了了解vLLM和其他开源项目到底谁快谁慢。这个benchmark可以在每次大版本更新的时候发布，但是要尽可能的详尽，易于理解。
于是，我说服了Ion，开始着手设计两套性能基准测试。
对于vLLM自己和自己比的测试做起来相对顺利一些，最终做成了一个网站。网站展示了与一天前和一周前的 vLLM 相比的性能变化百分比，还有折线图帮助定位是哪个 commit 导致性能变化。
第二个测试是为了在vLLM和其他项目（LMDeploy, TGI, SGLang, TensorRT-LLM）之间进行尽可能公平的比较。在开始深入研究之后，才发现这个方向比我想象中的麻烦很多很多。
首先是把测试跑起来就就不容易，主要的复杂度来源于TensorRT-LLM。 TensorRT-LLM当时在业界公认很快也公认很难用。我参考NeuralMagic的一个demo版本的实现磕磕绊绊的跑起来，但是性能却不好。最后直到和NVIDIA的人联系上了之后才跑起来一个相对正确的版本（这个版本都和NVIDIA的内部实现比还是更慢，但我们没找出为什么）。
同时，为了保证测试只需要复制粘贴加回车就能自动跑起来，我把这个benchmark做到了自动测试里面。我无数次发现一行代码不对，改了一下，然后就花半个小时到一个小时等自动测试跑起来，然后跑了一会崩了研究为什么，然后再改一行代码，再花半小时，如是反复。
测试总算能顺利跑出结果之后，随之而来的又是有关公平性的一系列挑战：
- 没有办法限制输出长度，就算测试的是一模一样的模型，由于实现的不同，用的库的版本不同，采样方式的不同，输出长度都不一样（这个问题随着大家都实现了
--ignore-eos之后算是有了一个答案—但这已经是很久之后了）。我调了很久TGI的输出还是莫名奇妙的不一样最后只能选择不测TGI。 - 看似一样的参数语义是不一样的（比如GPU utilization，有的是整个GPU的utilization，有的是刨除模型之后剩下GPU的utilization，有的项目会预留出来临时GPU变量的空间有的不会，很难做到彻底的对齐）。在发现这个问题之后我已经没有足够的时间来继续修复了。
最终赶在 vLLM 团队发布性能优化博文时完成了第一版。事后收到 NVIDIA 和 SGLang 的反馈，发现性能比较还不够公平。第一版基准测试确实不够完善，但至少实现了一键运行的目标：我相信只有让更多人运行基准测试，才能发现并解决问题。
研究方向：分布式推理，PD分离与KV cache传输
在performance benchmark告一段落之后，我对vLLM这个软件也更加了解，开始能够偶尔蹦出来一些研究想法。但是人只有一双手，我也只能集中精力攻克一个研究，那到底应该选择什么研究方向呢？
说来有点好笑，明明知道研究最重要的就是先做起来，我却举棋不定。
一开始我觉得，可能是我把这次研究看的太重了：这估计是我毕业之前能做的最后一段研究，这段研究也是我真正能横跨学界和工业界的一段研究（感谢vLLM让我能够接触到真正的LLM从业者），我也希望这个研究既能让我由衷的喜欢，又能获得工业界的影响力。
后来听了 Ion 的一次演讲，他说区分顶尖研究者和最顶尖研究者的，不是解决问题的能力（这个大家都有），而是选择问题的能力。这让我开始认真思考如何选择研究问题。
最终我选择了做KV cache传输这个方向，原因有四：
- 大势所趋：LLM越发多种多样。处理多种多样的LLM的最好方法就是把不一样的功能放在不一样的机器上，然后在机器之间传递LLM请求的中间状态（也就是KV cache）。
- 工业应用：工业界已经在开始实现PD分离技术，KV cache传输正是PD分离技术的核心。
- 技术能力：我之前的研究方向是视频传输，虽说视频传输和KV cache传输肯定是不一样的，但是抽象来说大家都是在传tensor。我相信之前的研究能帮到我。
- 个人口味：相比于建立一个复杂的系统，我更倾向建立多个简单的系统，然后优化系统间的传输和沟通来解决问题（从这点来讲虽然我在做MLSys但是内心还是个计算机网络研究者）。
带着这个念头，我开始借由PD分离这个具体的应用场景下写第一版KV cache传输的代码。意料之外（又在情理之中）的是，之前已经有很多版的实现，但都最终没有了下文。表面看来是因为之前的实现带来几十个文件的几千行代码改动，但根本原因也许是大家思维上的惯性—认为PD分离和其他的分布式推理（比如向量并行）是一样的，要进行一样的实现，要由vLLM自己来进行PD分离以及一系列的管理。这产生了两个冲突：
- vLLM 的理念是优化单个模型实例的推理，而之前的方案试图让 vLLM 同时管理 Prefill 和 Decode 的模型实例。
- PD 分离强调 Prefill 和 Decode 的解耦，而之前的方案试图在 vLLM 内部同时实现两者，不可避免带来耦合。
而我的实现思想，是把Prefill和Decode放在不同的vLLM实例，然后专注于实现vLLM实例之间的互联。在这种思想的帮助下，最终我的实现对vLLM的核心代码改动仅有不到100行。
说来简单，这个实现前前后后花了四个多月。在这段时间我反复和其他人讨论。影响我实现的最重要转折点，是一个来自于我博士课题组的一个项目，叫做LMCache。这个项目专注于为vLLM扩展KV cache的存储空间，同时也实现了不同vLLM实例之间的KV cache互联。
这帮我补上了我的研究缺少的一块重要拼图。我的研究目的是优化KV cache传输，而优化传输的具体效果需要这个传输功能本身的实现足够优秀才能真正被展示出来.而LMCache（以及最近新开源的Mooncake）正好在优化传输功能。
也是在这个时刻，我意识到传输功能本身应该交由第三方实现，而我能做的就是为第三方实现提供一个规范的接口，同时把这个接口背后vLLM这端应该做的工作做好。至于这个KV cache传输的接口应该长什么样，又怎么和vLLM本身融合呢，我在和大家广泛的讨论（感谢vLLM团队，Mooncake，SGLang，NVIDIA在这个过程中提供的宝贵意见），希望最终能够交出一个让我满意，也能让大家满意的答卷。我们拭目以待。
题外话：为什么要做开源的研究
世界上99%的博士们是不会花大力气在做开源上的。那为什么Ion的组（以及vLLM团队的大家）在明知做开源项目所需的努力大部分都和论文没啥关系的时候，还是选择做开源项目呢？
我不清楚Ion是怎么想的，这里只是我个人的看法：开源牺牲了发论文的数量，换取了做出顶尖研究的机会。
系统研究是被真实世界的应用驱动的研究，所以顶尖的系统研究要夹在工程和科研的中间。只有有了工程的视角，才能拥有顶尖的技术，才能第一时间感知到现在工业界里面什么东西重要；只有带着科研的视角，才能看到什么问题有科研的价值，什么问题的解法背后有根本性的观念转变，而不是满足于只是把问题解决。而开源是最好的带给研究者工程视角的方式。只有做开源项目让很多人用起来，研究者才能真正融入工业界的社群，才会有人愿意帮你，告诉你什么问题是重要的，帮你解决工程挑战。虽然开源一定会牵扯很多的精力，但是可能没有更好的办法。
结语
时势造英雄，还是英雄造时势？我想，是英雄与时势的共舞。
我深信大语言模型会变成搜索引擎一般的基础设施。在信息时代下，信息只有经过处理，才能变成能被人吸收的知识。而大语言模型会是连通信息与知识的桥梁。
大语言模型提供了如此多的机会，离不开其创造者们不计成本的付出。在这个浪潮之下，又将涌现出一个又一个新英雄。
大语言模型时代，我们都是造梦者。但梦造着造着，或许就变成现实了呢？