清华陈文光:用Claude Code开发RTX 5090注意力算子,AI 正在重塑底层编程的边界
AI 编程技术已经进化到什么程度?清华大学计算机系教授已经用 Claude Code 来开发高性能GPU算子,并带来极大的效率提升!
近日,在端侧智能北京市重点实验室主办、面壁智能及 OpenBMB 开源社区联合主办的「智能体创新研讨会」中,清华大学计算机系陈文光教授带来了一场别开生面的演讲。
他以一个硬核的工程实践——「在RTX 5090消费级显卡上,利用FP8特性开发一个比现有方案快20%的Flash Attention算子」为案例,详细拆解了他与Claude Code在过去一个月里「结对编程」的真实经历。
这不仅仅是一个追求性能极致优化的技术故事,更是一场关于未来软件开发范式的深度探索。陈文光教授指出,技术发展正从ChatGPT这类停留在信息空间里的对话工具,到Claude Code这类能真正操作计算机、调用工具、执行复杂任务的智能体工具,软件开发的方式正在发生根本性变革。我们正从「AI辅助编程」迈向「AI主导编程」的新时代。
这也引发了一个行业思考:在 AI 时代,程序员的核心技能是什么?还需要死磕C++语法细节吗?还需要懂汇编吗?——陈文光教授以一个「问题」而非「答案」结束他的演讲。以下是其演讲内容的整理精编:
何以引入Claude Code?
我今天分享的内容和学术没有任何关系,基本上是一个工程内容。我们想做一个稍微复杂一点的程序,看看现有的人工智能到底能帮我们改进到什么程度,于是大概花了三个星期在一个低端消费显卡上做了一个GPU的kernel,使其工作效率提升了约20%。
刚才知远(清华大学计算机系刘知远教授)提到,他觉得OpenClaw的出现,其影响力不比ChatGPT小。我的看法是接近的。ChatGPT这类对话模型,基本都停留在信息空间里,查消息、回答、生成报告等。但从Claude Code开始,能力就变成了使用工具,特别是操作机器的能力大幅提升。
特别是到去年12月,Anthropic发布了新的Claude模型,以及后面OpenAI发了新的模型,我认为整个软件开发的方式发生了巨大的变化。过去大家还是用Cursor辅助编程,但在我这一段开发算法的实践中,已经不是AI辅助编程了,基本上是用自然语言跟它对话、较为精确地提出需求,Claude Code基本上都能实现,而且实现的效率远远比个人要好很多。
我觉得Claude Code这类工具,通过很强的模型能力以及工具调用能力操作你的电脑:它能通过你的各种授权,用shell脚本,操作你的浏览器,可以做非常多的事情。它不仅仅是在对话框给你展示结果,而是真正能发邮件、去做好各种定时的事情。
而OpenClaw在这个基础上,把用户界面又做了大幅度的简化。Claude Code还是针对我这样的专业人士——平时在电脑前工作,写程序、写文档、发邮件等等。但OpenClaw就把接口彻底打开了,它可以通过即时通讯软件访问AI。底层逻辑是一样的:选有用的工具去调用。任务可能非常复杂,需要多步的工具调用,但是它们共享同样的核心。最近人工智能的进展使得这样的结合变得可行,这也是为什么OpenClaw在今年1、2月开始变得非常流行。
但OpenClaw和Claude Code还是有一些区别的。Claude Code更多的是一个项目的概念,源文件会放到项目的目录,它首先会搜索文件夹里的代码、文档;OpenClaw的界面则更像是一个功能集合,我今天和你说五件事、八件事,它能很好地组织起来。AI有一个很重要的上下文或记忆的问题,OpenClaw需要做更好的管理。
总的来说,我对于OpenClaw还是一个浅度用户,所以我后面更多的分享是关于Claude Code的。
用Claude Code开发高性能GPU算子:我们所做的事情,概括来说就是:在消费级的RTX 5090显卡上开发一个算子,叫作Flash Attention算子。:为什么单独干这件事?如果大家关注AI芯片,就会知道:英伟达的生态很好,有CUDA兼容性,比国内好。但如果仔细了解就会发现,每代英伟达架构的差异都非常大,因此,所谓的「CUDA兼容性」其实也没有那么兼容。换了一代新的架构以后,基本上所有的高性能算子都要重写。目前很多程序员、很多钱都花在怎么在一个新的结构上写一个更快的算子。:英伟达的RTX 5090比H系列还要再往后一代,更先进,但它是用来做消费的,虽然也是B系列,但和同为B系列的B200、B300这些卡,在底层的体系结构有非常大的差别。但它们可以共享一些东西,比如都支持Blackwell FP8,支持一种新的数据格式:一堆参数共享一个指数。这种格式在原来的H系列上是不支持的。:所以5090和B200都有我们认为在训练推理上有很大帮助的数据格式。实际上,B200确实把它发挥出来了。如果大家关注的话,Flash Attention 4 最近刚发布,它使用了很多新技术,包括把FP8用得很好,得到了很好的效果,但5090并没有匹配的算子,这块并没有写好。
现在Transformer Engine也实现了一些在5090上的功能,它是在CUDA NN之上再包装一层,方便大家做AI训练。它支持了MX FP8的训练策略,其中细分了几部分:一些线性层用了MX FP8,但注意力这个部分还是用的BF16。
我这个项目就想说:既然硬件有这个能力,现有的库又没有提供,我能不能利用这个硬件提供的FP8能力,实现一个更快的Flash Attention?这样我们就可以用5090这种便宜的消费显卡做模型训练,给模型训练提供一个低成本的解决方案。
实际上这个事情非常复杂。5090看上去是B系列的卡,但它和B200的结构非常不一样。首先,它没有Tensor Memory,也没有另外一种异步计算的指令,甚至最重要的访存和计算能力都不足。只能每次把数据从整个卡上的HBM(全局内存)加载到计算单元可以快速访问的共享内存(Shared Memory)上,再加载到寄存器里进行计算。
这里就发现一个问题:这个Shared Memory其实非常小,每个Stream Processor只有128KB,这对算法编写会有很大的影响。
正是因为5090没有Tensor Memory,所以在B200上面那些优化的、基于Tensor Memory的方法都不能用。需要计算怎么办?就必须把它从Shared Memory再加载到寄存器里,在寄存器里面来做FP8的MMA(矩阵乘法累加)计算。这时候寄存器容量又成为问题,每个线程只有255个寄存器。如果使用超出,就会有所谓的“spill”,把当前不用的写到Shared Memory里面去,需要时再读回来,这样就会慢很多。
由于体系结构的限制——没有Tensor Memory,没有所谓的WGMMA指令,访存和计算需要比较死板地自己写,再加上Shared Memory和Register都非常小,这给我们后续写算子带来了巨大的障碍。
实际上,Attention有前向和反向。前向就是顺着Attention架构算回去,反向是算它的导数后再计算回来。前向两三天就做完了,而且性能很好:FP8的前向是BF16的前向性能的1.7倍,这已经相当不错了,而且寄存器使用也很好,没有spill。但反向遇到了非常大的问题,因为反向有五个矩阵乘,它们怎么互相连起来,怎么传数据,同时保证Shared Memory和寄存器不超出,这花了我们非常多的力气。
这个工作不像一般的前端程序,它基本上是一个逻辑流的过程,中间有很多步决策和优化:我要往这个方向走还是那个方向走?先探索什么方向,后探索什么方向?而且任何一个优化,都有可能提升计算效率的同时增加了寄存器压力,导致最后实际上并没有性能提升。
这是我们性能优化的一个过程。下面这个图是Kernel的执行时间,从一开始八点几个毫秒一直到后面的三点几个毫秒,下降了很多。右边是和BF16的对比,在TransformerEngine里面实现的工作,最高到过它的1.25倍。但是这个1.25实现的版本不太对,有BUG,修好BUG的时候在某些形状上可以达到1.2倍的加速比。
最有效的优化,在整个过程里面探索出来的有几点:
第一个是在「计算完再去做乘法」这个环节,两者时间都耗着,你应该做一个迭代的,这边做Load,那边做计算。另外还做了一个「unroll」,如果把它展开以后,分支会少一些,融合的情况会好一些,一般性能也会增加,但同时也增加了寄存器压力。正如我们之前所说,寄存器非常有限,所以unroll 1的意思就是不unroll,降低寄存器压力,把spill从一千多byte(字节)降到272个byte,性能提升很多。
第二个优化是做了一个矩阵的in-place的transpose,在Shared Memory里面做,这样就避免了要把它写回到global memory再读回来——因为Shared Memory里边没地儿,不做这种in-place transpose的话,那就没地方放,只能倒出去再倒回来。
这两个优化效果都很好。
其实这个AI不是傻傻的,并非你让它干一件事、它就只干一件事。我说要优化,它就会自动提出SASS框架,分析你的指令、每类指令的分布、以及它应该达到什么样的性能。AI有一套系统的方法论,让它优化性能,甚至不需要我来告诉它。但是我还没有给它提出用Nsight那种需要超级用户权限的东西,这种在用户态做的分析或许还能得到更高的性能。
细节的工作我就不展开太多了。我大概讲一讲,在我跟Claude Code一起工作了一段时间以后,什么事情是AI能干好的呢?
第一个是体系结构的探索。它可以去分析现有的代码,哪些指令有多少、有多少个spill,包括读数据的时候要把它散开,才能把性能发挥出来。这些事情我们人去做很累、很机械,但AI做得很好,不知疲倦。这些对它来说是很例行的事情,基本都能做对,不会犯很多错误。
第二个是系统性优化。我们人要找到很多热点,再看看是计算有冗余,还是访存有冗余,还是计算和访存重叠得不好,还是多个线程负载不平衡等等。但这些事情,Claude Code都能相对自动化地去实现。在这个过程中等于内嵌了一个性能工程师的基本能力,这应该是从它的训练数据里来的。每次我们都会做正确性测试,对的我们留下来,不对的就给它反馈回去,这样的迭代就能做得很好。
另外像“pingpong”跟FA4这种把线程分成两类,一类专门用来计算、一类专门用来访存的方法,我们也试了4+4、6+6、8+8各种方式。让我自己写,早就写死了。但让AI写,不对就撤回去,对它来说就是尝试,失败也没有关系,机器是不知道累的。包括unroll与否和一些参数空间的探索,都可以做。因为我们是在Flash Attention上改的,要和底层的像cuBLAS、我们的训练脚本这些做一些集成,这些事情系统也能相对自动地做得比较好。
另外,像我们做5090(SM120)架构的算子,在我的提示下AI都可以读Hopper架构(SM90)以及 B200架构(SM100)相关的论文、报告,来看有哪些技术是可以用到我们这里来的。这些事情。包括最后生成一些报告,对它来说都是比较容易的事情。
最后也稍微介绍一下我们没有探索成功的,但这其实不太重要,反正我们自己工作也会有失败的方向。
一个方向就是4+4 pingpong,做之前我们还做了一个估算,到底可以提升多少?我们让计算和访存重叠起来,能够得到一个比较好的性能。但是真正用到实际的完整的Attention的时候,寄存器压力等等会受到很多限制,最后做不出来了。
但也不是说所有探索都是没有用的,实际上这些探索逐步地让你对这个问题有了更深的理解。在8 warps、128-byte rows、255个寄存器的占用下,这是我们最后得到的一个稳定结果,因为做别的尝试总会破坏其中的某一些因素,让整体的性能变得更差。
未来,AI 时代的编程会变成什么?
我一边做一边想:如果换一个人来做这件事情,能不能做得更好?或者说哪些方面可以做得更好?
第一个,就是我刚才说的,没有把Nsight那些底层的细节开放给它。这不是系统性的大问题,是细节问题。
第二个其实很重要。需要管理内存、管理记忆。其实对复杂的任务来说,要读很多代码,搜很多东西,还要读很多实验的结果,所以Context Window我用了一兆上下文的模型,但做着做着,到后期,还是会出现一些情况。还是要时刻关注这件事,给它做层次化的内存组织,时不时地做一些压缩让它清空。
一些比较常见的技能,大家熟悉OpenClaw的话都已经知道怎么使用了,但plan mode这种复杂的事情,不能让它做太多太冗杂的事情。一个难的事情,先让它思考一下,定一个详细的计划,后面再去实现那个计划。其实系统上也有一个brainstorm的skill,装上去之后效果会更好一些。
在使用过程中我们常遇到的认为能减少寄存器压力,但实际上并没有真正减少的这个问题,究其原因英伟达只开放了PTX指令集,这是一个虚拟指令集,它对于下面那层的寄存器分配其实是没有任何控制的,负责寄存器分配的编译器对AI来说还是一个黑盒。寄存器分配这个部分要通过外面程序去改进它,其实很难,因为你也不知道这个寄存器分配算法本身是不是已经把临时的用不到的变量替掉了。
这两个方面的配合,目前模型还没有很好的办法去完善。我如果只改一个经常改的文件,大概需要一分钟左右的编译时间;如果我改的东西多了,差不多就需要一小时的时间。我开始的测试相对比较简单,到后面发现loss(损失)是有问题的,后面都用比较长的测试,长的测试大概十到二十分钟的时间。
其他也有一些失败的例子,比如AI没有得到许可就直接修改导致错误,也是有这样的情况。但是我们后面还是做了非常好的软件工程的保证,每次修改以后都会做相对完整的正确性测试和功能测试,有好处的留下,没有好处的就反馈回去,整个版本管理让它做得更加安全一些。
最后,虽然有的张量形状我们实现的算子比过去BF16的Flash Attention快了20%,但引入到模型训练测试的时候,实际训练速度是相当的,并没有快很多。能不能有机会把FP8用好,真正把它的潜力发挥起来,还需要有创造力的一些工作。目前的AI水平是找一找别人干了什么、看一看能不能用到我这个场景,真正完整的创新还没办法做。但我并不觉得AI不能做,本质上是搜索、强化学习和人协同的工作,但至少目前还没有做这个事情。
算子做完了以后,性能还可以,我就发给同学试了。但很快就发现不太对,这其中有一个收敛的问题,到几十个迭代,到60以后有一个问题就比较明显。所以我们那个学生觉得有一些错误,自己还写了一些测试的脚本。
我这两天在修这个东西,一直用各种各样的方法,也发现了一些BUG,也把它修好了。基本上迭代到80的时候,之前收敛到5.96,BF16应该到5.5左右,所以现在在80个迭代的时候也可以收敛到5.5左右,基本上把它做好了。
对GPU kernel来说,我一直在做国产平台性能训练。过去有一个看法,国产平台的主要问题是没有高性能的kernel,所以如何给国产平台写高性能kernel是最大的问题。
我做了这个事情以后,我发现英伟达平台其实也并不简单。看上去CUDA兼容,其实底下的东西差的非常多,很多东西不能复用,基本上要重写,因此出现了很多新的编程抽象,像Triton,像北大最近做的TileLang。
大家都希望说,我能够有一层对程序员更加友好的抽象,同时还能控制底层,再加一个很通用的编译器,来把底下的算子做好。我经过这段时间的实践以后,我现在稍微有点怀疑这条路,因为非常底层的东西,像5090上面的很多体系结构特点,都没有支持。底层的Triton编译器、Cutlass等都没有对5090提供很好的支持,要开发高性能算子只能写最底层的CUDA,甚至要写PTX汇编,才能把它的性能挖掘出来。
如果是这样的话,这种中间层抽象的价值到底何在?我现在稍微有一点疑问。但我这个工作还有点初步,所以这不是一个非常确定性的看法。
对于我们用AI做各种各样的工作,实际上是因为你要做不同的任务才需要考虑到内存的使用,甚至我们在这一个任务里面要分很多的子任务,多少成功多少失败,每一次都是大的代码重构。怎么把内存记忆管理好,能够把需要的东西放到上下文里面去?现有的「skill机制」和「memory机制」,最基础的逻辑还是做层次化记忆,一次不要都堵进去,我是按需地把以前知道的放进去。目前来说,在整个管理上还没有做到完全的无感知、自动化的把它做好。
所以我认为越懂这个事情的人,依旧越能够把自己的能力发挥出来,做得更好。编译优化和性能优化里面的基本概念我是很熟的,对于CUDA的编程基础我也知道。但是对于5090这样一个我还不大清楚的细节的体系结构,我还是能够和AI一起做到一定的程度。
这是整个项目的一些情况:大概写了5000行的CUDA/C++代码,改了30个文件,做了大概19个成功的优化,15个不成功的优化,也支持多种GQA形状。前向最多做到347 TFLOPS,反向做到240多。大致上,AI可以帮助系统专家放大自己的能力。我们懂的东西能用AI扩大你的能力,可以做得更好,但不是完全的替代。
AI 时代,我们还需要学C++吗?
最后我想用微博的一个帖子作结:
于仕琪老师是南方科大的一个老师,他非常负责任,他是讲C++的,他说C++将进行古法Coding:把面试纳入内核,因为留作业已经没有意义了,不管留什么作业,学生交给模型做肯定都能做出来,况且作业比我今天讲的这个东西简单一万倍,随便就可以做出来,没有办法靠作业评估学生到底掌握了多少。
后面第一个回复是我。我当时想,学生到底需要学什么东西呢?我觉得他如果真的需要学,你怎么考他都没有问题。但是现在这个情况我稍微有点不自信。因为现在很多人对汇编已经不熟悉了,包括我们去写操作系统的人,其实也不是那么熟悉了。现在的高级语言以后会不会像汇编一样成为极少数人的需求?也就是我们只需要用自然语言和模型对话,模型产生高级语言再靠编译器去编译,也可以做很多事情。但这个事情还是一个问号。
这条微博的另外两位转发者在国内开发界也很有影响。一个是说如果汇编没有用,就别开课了,开课就好好教。邹欣说培养CS核心技术的科研人员和培养应用人才课程的要求大不一样。我觉得他目前的认识可能在跟我在半年前的认识差不多,觉得AI只能做一些个人用的一次性网站之类的。经过这段时间的实践,对他目前的说法稍微有点怀疑。难道写kernel的不算CS核心的科研人员吗?我觉得也算吧,但是好像不懂底下的细节好像也能够做很多的事情。
所以右边这个问题,AI时代,CS核心技术的科研人员需要知道C++语法细节吗? 我其实不是很确定。我大致感受到的一个方向是,和AI的协作、宏观的知识更加重要——你要求它做什么,以及往哪个方向去优化。
你可以想象我们在一个大的搜索空间里面剪枝,让AI少花力气。因为无穷投入干一个事情,投入大于事情的价值是不行的,要让投入小于你产生东西的价值这样才有意思。所以剪枝还是很重要的。
微观的知识我也不能说一定不需要,包括汇编、写编译的人还是需要的。可能我觉得类似于CSAPP那个课,当时也有人问我,大家上《深入理解计算机系统》这门课有什么用?那个课里面说,不一定自己能写汇编,但是要读懂汇编,知道编译器会生成什么东西,知道你写一门高级语言,编译器会把它编译成什么样子。我们把它对称回来,就是:你也许知道一段自然语言,模型就会把它实现成什么样的代码。