Hongzheng Chen Blog
PhD第三年 - 漂泊与归途
Dec 31st, 2024 02024，我在美国读博的第三个年头结束了。
今年的主题是漂泊，坐了31趟飞机，走过了30座城市，开车开了快两万公里路，大部分时间都在路上。发生了太多太多事情，以至于这篇总结拖了几个月才开始写。去年还一篇paper没有，今年存稿的5篇就全中完了，顺便拿了些奖，给了不少talk，拿了master学位，也拿到了NVIDIA的全职offer，似乎差不多可以毕业了。
但光鲜亮丽并不能概括这一年。25岁，也是我极端动荡的一年，生活上经历了各式各样的痛苦和磨难，焦虑与无助在这一年如影随形。quarter-life crisis，在奔三的年纪变得更加具象化。尽管这一年过得很艰辛，归途是什么还没有寻找到答案，但在一遍遍跟自己的对话中，在如同家人般的朋友们的陪伴下，还是收获了些成长，和继续前行的勇气✨
理想的生活
这一年绕着地球飞了两圈半，去了很多我以前一直想去没机会去的地方。想不到自己竟然慢慢当上了“数字游民”，这是我十年前就幻想过的生活。这一年去的30座城市，每一座现在回想起来都仍历历在目，他们过于棱角分明，以至于我现在还是能清晰地回想起在每座城市短暂几天时间里感受到的声音、温度、还有风的气息。
今年平均每两周就要坐一趟飞机，总的旅途相当于绕了地球两周半
这一年的旅途是关于勇气。在Killington的山顶面对近乎跳崖的双黑道，第一次敢自己一个人往下冲，没有叫救援一路滑到底已经是最大的突破；在San Diego体验了滑翔伞，冲出悬崖在天上飞翔的体验真的太特别了，看到了最温和的加州的海，还追到了一群白色海豚在海面上嬉戏。
是关于跨越经年书中与现实的重逢。高中毕业旅行之后就再也没有去过日本，年初趁着转机的时间回到东京，赶到东京塔下时铁塔闪耀着橘红色的光芒，星星点点的雪花正好飘在身上，好像记忆中的一切又都回来了；去了心心念念的Key West，那是美国的最南边，连串的岛屿连绵入海，小镇不大也没有太多的景点，但是在那我看到了十年前在书中描绘的果冻海的模样；还有以前人们一提到加州/洛杉矶必定会提起的Santa Monica，虽然没坐上落日飞车，但是夕阳下游乐园小朋友们的笑声此起彼伏，让那个下午变得更加难忘。
Lost in Tokyo
是关于人与人之间的奇遇。在哥本哈根餐厅偶遇了个波兰小哥，在诺大的落地窗前看人来人往，我们讲着各自国家的故事，短短聊了两个小时却仿佛已经让我读完了他的人生；在西雅图跟认识不久的日本朋友彻夜长谈，第一次发觉人与人之间的共通并不会因为国籍性别还有所处环境而受到阻碍。
也关于工作。试过一大早起来从村里开车一路狂飙穿过加拿大边境，到Montréal直接setup好跟老板开会；也试过飞机落地之后直接跟人约coffee chat；在滑雪场回Slack，在车上听talk，在飞机上开会这种事情今年都干了不少。
今年因为中了几篇paper，所以得以参加了不少conference，在会场上能够被人一眼认出来，即便很多人以前从来都没有见过。越来越多人知道我不是因为我的老板或是仅仅因为我发了paper，而是因为我做的工作，因为真诚地觉得我的方法能够解决实际问题。收到了很多的鼓励和称赞，我非常感激。有些素未谋面的人也因为读了我的博客，专门写邮件来感谢我，说这促使他们想要读PhD。其实对于我来说也非常感动，我从来没有想过以前的一些随笔能够潜移默化地给别人带来影响，如果短短一些文字能够激励更多的人去追寻他们所喜欢做的东西，那还是件很有成就感的事情。这一年发现能够创造impact的东西不是只有论文和citation，还有其他很多很多的东西，包括博客、摄影、视频、所有跟创作相关的其实都能够创造impact去影响更多的人。也许记录的目的就是为了记录下那些鲜活的瞬间，为了说明PhD不是都要“甘坐冷板凳”或者卷生卷死，每个人都能够有不一样的活法和生存方式。
破碎的眼前
前面的诗和远方很多时候仅仅是表象，光鲜的背后是更多的痛苦和挣扎。对于我来说，科研或许是件简单的事，怎么好好去生活反倒是更困难的事情。
从来没有人跟我说过真正成年的生活这么艰难，去年的总结里以为自己拥有了稳定的内核，已经有能力承受一切不确定性，到头来还是被一件件事情击得粉碎。去年底回国前出了起小车祸，后来车借给朋友开又被撞了两次，单是这一年车就修了三次。虽然人都没有事，但是频繁跟修车和保险公司扯皮非常耗费心力。
而更加黑天鹅的事件是被告知家里人重病，一下子承担了家里最重要的责任。突然间发现死亡离我这么近，对未来的巨大恐慌导致的压力和焦虑，进而诱发躯体化症状跑了不少次医院。墨菲定律告诉我们坏事只会接连发生，太过想逃离糟糕的状态，反而会适得其反。滑雪扭伤手腕连续几个月都没有办法正常用惯用手，人际关系上也是莫名其妙被不同人中伤拉黑，一件件事情仿佛压在我身上的一根根稻草，到某一刻就会突然崩溃。
而当我还在迷茫的时候，同龄人都进入到了人生的下一个阶段，毕业回国、工作、亦或是结婚，似乎只有我还在流浪。在还没出国前，我几乎百分百确定我会在美国拿到学位后立马回国，但在美国呆了几年发现自己不想回去了，而且也回不去了。在美国的几年时间仿佛被压缩了，两三年时间经历的事情比我过去十几年经历的事情还要多。今年一直在漂泊，我妈说累了就回来吧，差点没忍住爆哭。即便我回去了，大概也回不去以前那种无忧无虑的日子了。
一直以来都觉得我的运气很背，好事情不会发生在我的头上。比如说今年的几大天象日全食、极光、流星雨我所在的地方都是多云；让老板帮忙推了三四个fellowship没有一个最终拿到，尽管我觉得已经将所有我能做的都做了。今年真正的turning point也许是决定回国办签证。本来在办签证前两周还犹犹豫豫左右摇摆，害怕被check或者被拒回不了美国没办法继续读PhD。回国前一晚在公寓楼下跟王总聊了半个钟，最终被王总劝服莽一把。大概是发现很多事情如果不去试一试那会永远后悔的，努力过至少不会后悔选择去做。
面签那天早上六点多就到北京大使馆门口排队，在心中演练了无数种结局，已经做好99%会被check的准备（之前签了几次都是被check），但面签实际两分钟，VO没有看什么材料，问了一两个问题就直接给我过了。我确实没想到会直接拿到签证通过的蓝单子，我拿着单子一路蹦出大使馆，那或许是我这一年最快乐的时候。
面签完之后跟1+在咖啡厅等王总出来有种像在等孩子放学回家，最后是皆大欢喜，感谢王总和1+！
虽然痛苦是生活的必然1，但是我们还是要去寻找和感知那一个个鲜活的瞬间，这些片刻是生活的疗愈剂。这一年认识了很多人，也看到了很多25-30岁之间朋友的活法，他们的自由自在不被现实所束缚的天性是我所向往的，他们用亲身经历告诉我年龄增长、面对那些生活的变故似乎也不是那么可怕。生活从来就不是一个简单的线性问题，很多时候甚至不是well-defined的，每个解都会有自己的tradeoff。其实人生中很多问题更像是在摸彩票，也许这边运气背了点，别的事情运气就会好起来了呢。
这一年被迫去成长，被生活中很多突如其来的事情推着前行。很喜欢我前主管说的一句话“抱最好的希望，做最坏的打算”，既然都已经努力了，那就不要再责怪自己了。这一年最感谢的还是自己，我居然熬过来了！虽然自我重建花费了很长时间，但那些荆棘与痛苦没有把我磨砺得麻木不仁，我对生活的感受力还在，我依然没有对这个世界丧失了热爱。
科研的进展
而科研，大概是在我低谷的时候依然能够让我燃起信心想要继续做的东西。
Allo这个工作从入学就开始做，做了三年时间终于中稿了。其中的故事也非常曲折，从最开始入学时接手HeteroCL2，将整个基于TVM的后端重新改写成基于MLIR的框架，因为仅仅是将decoupling的想法在IR层面上实现，我觉得从novelty角度上来说并没有什么新颖的东西，更多还是engineering。第一年想了很久没有找到合适的storyline，所以就一直拖着没有动笔写paper。第一年结束后我去Amazon实习，正好赶上ChatGPT出现，于是中间一整年时间都在做分布式并行训练框架Slapo，而把学校的项目搁置了。等到第二年结束，有了加速LLM的一些经验，突然感觉我们可以基于这个编译器去搭一些大一点的东西，而LLM正是一个绝佳的例子，所以我们就用HeteroCL-MLIR开始搭建LLM的加速器。去年暑假来了几个intern，在搭建加速器的过程中遇到了很多问题，主要原因还是基于TVM frontend的编程框架非常难用（HeteroCL用的还是古早v0.x的版本），所以后来我们又以此为契机重新设计了整个frontend的编程语言，并命名为Allo，然后才有了这整个编译栈。
最后准备投稿PLDI的时候只剩10天时间，最开始还没想好第二个部分关于composability要怎么描述，只有一个比较general的问题是怎么将有不同interface的kernel组合到一起。因为MLIR将layout也作为memref的一个type表示，所以如果layout不同导致interface不一致，在IR层面生成的module就已经无法通过verification。有一个比较直接的想法是将不同layout显示地用type描述出来，后面将kernel之间的dependency formulate成一个graph，interface unification的过程实际上就是type inference的过程，而将在这个dependency graph上做layout propagation实际上又等价于去做static dataflow analysis。要让type inference变得高效而不用SMT solver的方式是用Worklist algorithm，如果能够证明这些基本的partition type构成一个lattice，并且数据流的转换关系是单调递增的，那么基于不动点理论，这个算法就能够在线性时间收敛。
花了几天时间将这些都连起来想通的时候真的感觉Eureka了，这也是做这几个工作做的最有意义也最快乐的一次，最后写成paper没有经过recycle就直接中稿了。在PLDI的论文里，我们其实只讲了怎么处理memory partition的情况，但很多layout相关的问题都可以被归类到这个框架里解决。
今年做的另一个工作则是跟Allo同期的工作，我们尝试将LLM部署到FPGA上。这本来是一个课程项目，只基于BERT做了些实验。有了Slapo的一系列基础之后，我在这个工作里把GPT给加上了，并分析生成式推理（generative inference）在FPGA上做的可能性。本质上这篇论文想论证的就是我之前说的暴论——FPGA除了AI啥都可以做，但就是不能做AI。这篇论文发表之后有了很多follow-up的工作，但是我做完这篇之后其实已经不太想再碰FPGA加速LLM的项目，毕竟现在的FPGA无论是算力、内存、带宽都全方位被GPU碾压，对于Transformer这种规整的模型就更加难有提升的空间。
因为很清楚用FPGA做LLM不是出路，只有用GPU才能做出真正有impact的工作，于是我去年底投了NVIDIA的实习。本来就抱着试试的心态，因为我的主业并不是GPU优化，也没有基于CUDA搭过大型项目，所以并没有想过我能真正拿到offer进到NVIDIA实习。后来居然面过了，也去到了一个专业对口的组做compiler。刚进去的时候我manager让我挑选项目，一个是分布式并行，另一个是底层的kernel优化。因为我去年已经在Amazon做过类似的项目并且熟悉各种分布式并行的策略，所以如果我做第一个项目其实没有任何入手难度。本着对新东西的热情，以及想做一些真正GPU相关的优化工作，我最后还是选择了第二个项目去做底层的kernel优化。
这个项目主要针对新的GPU架构进行编译优化。自从Hopper有了TMA和WGMMA这些异步模块和指令后，就可以做非常细粒度的流水线（pipeline）了，将不同的数据读取和计算的阶段给overlap起来做成数据流（dataflow），同时可以利用不同的warp实现不同的producer和consumer的功能。这样的好处在于可以边读边算，从而最大程度降低数据读取的开销。实际上FlashAttention3就是针对Hopper架构实现了一套warp specialization的算法，但它是人工实现的kernel，需要的人力成本非常高。现在的application太多了，attention也有很多变体，怎么对不同的这些kernel都能自动生成warp specialization的GPU code是我这个实习项目研究的问题。
这个项目对我一个只有本科生水准CUDA代码水平的人来说太有挑战性了——现在的GPU高性能程序已经复杂到了令人发指的地步，哪怕是一个简单的矩阵乘，在Hopper架构下的优化用CUDA写也要超过1000行代码。我对着高度优化过的CUDA代码啃了一个月，最终终于搞明白里面的每一个细节，然后用Triton实现了一个MLIR prototype，将自动做warp specialization的方案实现。一开始接到这个项目的时候，从没想过能够在暑期十几周内完成，毕竟这个项目一般是一整个团队做的事情，但没想到在暑假实习期间就把这个prototype搭出来了，并且跑了一些实验，在FP8 GEMM上面的一些case竟然可以比手写cuBLAS还要快1.5倍以上，而后续生成的FlashAttention kernel也能在Hopper上达到FA3的95%的性能。这个项目做着做着机密级别也越来越高，所以具体的实现方法在这也不能描述了，可能之后才会被发表出来。
而在我实习期间，NVDA的股票一涨再涨，一度跃升市值最高的公司。在NVIDIA的日子里，做大了影响力，被两个组的VP在邮件里讨论我的工作，申了专利，以内部价买到了股票，还拿到了全职offer，基本上把在NV能够拿到的东西都拿到了。虽然之后还没决定要不要入职NV，但至少我现在可以说我做的是全硬件（CPU/GPU/FPGA）优化了:)
未来的规划
今年正值我的高中百年校庆，2014年我们给十年后的自己写了一封信，今年正好到了拆信的时候。那时写的信还非常中二，想着要学很多高大上的理论包括泛函分析复变分析之流，结果从上大学开始就再也没有碰过数学了，而走向了另一条道路，现在的数学水平可能甚至还比不上当年高中的水平。但是很多年前畅想的去环游世界，去世界各地穿街走巷听不同的人讲他们的故事感受不同地方的烟火气，现在也还是喜欢并且认真践行了。也许漂泊从小就刻在骨子里了，过了这么久发现我依然是anywhere person，四海为家，总是在追求稳定和不稳定之间横跳。
从ChatGPT的诞生到现在已经过了两年，我们似乎离AGI近了一点，多模态的进展迅猛，能够生成更逼真流畅的视频了，也能让AI跟人类实时对话。LLM让很多事情变得简单，但也让很多研究变得没有意思了。Transformer横行的年代，很多人都只是将其作为一个工具去使用，去做一些prompt engineering，大多数paper看了abstract就已经知道有没有继续读的必要了。
对于system来说也是一样的，Horace在ASPLOS’24说真正有用的MLSys的工作就只有两个，一个Triton，一个FlashAttention，此言不虚。如果说两年前懂LLMSys的全世界还不到100人，那现在已经变成了显学了。怎么做KVCache又或者是更好的创建pipeline去overlap各个阶段，这些工作一开始可能还有点意思，做的人多了很多的idea也变成了排列组合，更多就是在拼手速了。在Amazon做完Slapo的工作后我本来想乘胜追击继续做分布式LLM推理相关的工作，然而去年暑假在肝了两个月做了些无用功之后，我决定quit。我发现与其去做一些太卷的方向，不如花多点时间去思考一些更fundamental的问题，这些给我带来的价值可能比citation大得多。
虽然跟很多人说过大概率会去工业界，但是随着现在工业界做的东西越来越同质化，很难说我会不会想去找教职。如果真正走上了教职这条路，也许我会选择去研究更编程语言（PL）相关的东西。我现在也会跟别人说我的主要研究方向是PL，毕竟PL魔法师3听上去比较酷一点，PL里那些深奥而无用的理论就像魔法一样或许能够减轻更多程序员的负担。如果能够像芙丽莲一样把“变出盛开鲜花的花田”的无用魔法教授给更多的人，给更多的人带去快乐，那我也会非常开心的。
最后以下面这句话作结吧
Per aspera ad astra
意思是『穿越逆境，直抵繁星』，这句话作为我的手机封面陪伴了我一整年。 希望在经历了九九八十一难后，最后一年的PhD旅途能够离那片星空更近一点。