摘要:近期,一位知友总结了自己最近三年在国产芯片硬件性能优化方向的工作经历。他接触了多家国产芯片公司的产品(代号 A、B、C、D),分享了在这些项目中遇到的各家芯片在硬件架构、软件编程及优化上的很具体的实操问题。有的太原始难用,有的学国际主流走得顺,有的想创新却没做
今日:国产芯片
近期,一位知友总结了自己最近三年在国产芯片硬件性能优化方向的工作经历。
他接触了多家国产芯片公司的产品(代号 A、B、C、D),分享了在这些项目中遇到的各家芯片在硬件架构、软件编程及优化上的很具体的实操问题。有的太原始难用,有的学国际主流走得顺,有的想创新却没做好,算是国产芯片开发现状的小小技术缩影。
在工作难题之外,他发现,国产芯片要想「快」,开发体验(手动管理数据搬运这种基础活都要工程师手搓)、生态被动(新技术都从英伟达生态诞生,国产只能追赶适配)等问题更亟待解决。
知友@有了琦琦的棍子
中国科学院大学 计算机系统结构硕士
7 月 28 日发布于知乎
刚好最近有时间,聊聊这三年在国产芯片上的工作经历。一方面是记录一下这三年做的一些事,一方面是基于这三年的工作经验引发的一些思考。 我是在 22 年年中的时候离开一家大公司去了一家小公司,主要是做高性能计算,针对各种硬件做性能优化。 这三年陆陆续续地接触了几款国产芯片,体验各异,挨个说一下。然后为了避免不必要的影响,用代号 ABCD 来表示。诸位且当听个故事。
先说 A 公司,当时做的是他们的第二代卡。 这个卡的 硬件架构可以理解成几个大核拼在一起,然后每个大核里面有若干 SIMD 单元以及矩阵计算单元。存储结构大概是 HBM、L2、L1、寄存器这几层。 硬件架构上我了解的少,不做过多评价。但是在我的视角上,芯片基本上都是这个样子,也没啥不好理解的。
至于在软件上,玩的就比较花了,并不是大多数芯片公司会采用的类 CUDA 的编程方式,更接近于 CPU 上玩 SIMD 的编程体验。但是比 SIMD 还难搞很多。
从访存角度来看, 这块卡的所有存储空间都是需要开发者来控制的。这意味着需要开发者来控制 L2 L1上放置的数据量,然后通过 DMA 操作来完成各级数据的搬运方式和搬运量。这跟正常的 GPU 有很大的区别,在绝大多数的卡上,一条「int a = A[0]」就能够完成从全局内存到寄存器的搬运。中间的 L2 cache 和 L1 cache 都是开发者无需关心的,由硬件自动来完成调度和控制,开发者需要做的也就是通过一些 hint,或者说 ptx 上的一些修饰符来提示硬件从而获得更好的访存性能表现。
而在这块卡上的一切数据搬运都需要开发者来手动完成,而且使用上全是异步拷贝的方式,你需要写一大堆代码来完成数据搬运的操作。而且为了提高性能,需要考虑各种因素,要通过 cache 的 size 大小来提前计算每次搬运的量,有的时候可能一次性搬运 1024 个数据性能差,但是搬运 2048 个数据就性能好。并且原生地支持搬运时候的 transpose 特性,又得考虑啥时候做 transpose 操作更好。全部都是异步拷贝,稍有不慎硬件就 hang 住了。
从计算角度来看, 软件层面上并没有提供 C 语言或者 C++ 之类的成熟编程方式能用。他们采用的大概思路是用一些 Intrinsic 函数来封装底层的指令,这种封装基本上是一一对应的状态。这可以理解成开发者只能通过汇编来写算子。而且由于计算单元是 SIMD 的模样,所以要想着在哪个层面进行向量化的问题,怎么处理边界情况。其中辛酸,很难说清。
然后访存和计算的问题又会交织在一起, 这里面体现出来的就是 需要极致的软流水操作才能将一个简单的算子性能实现的比较好 。举个例子,比如写个向量加法的操作。如果是 NV,基本上就是开些 block,一个线程负责几个数,然后进行累加,将结果写回。通过 GPU 里面的 SIMT 机制,让 warp切换来掩盖计算的开销,把访存单元给打满,这样就能够获得 95% 左右的硬件访存带宽了。但是在这块卡上,你需要通过大核数量,先确定每个大的核负责的区域。然后通过 cache 的大小来手动计算每次搬运的数据量,然后 L2 L1 甚至于寄存器都开双份空间,用来做 ping pong 操作。以 L1 为例,一半的空间是从 L2 上 load 数据,一半的空间是喂给寄存器,两者同时进行。再配上每一级的异步拷贝,从而让 DMA 一直处于 busy 的状态,这样才能把访存带宽打满。
这一套编程方式导致了在这个卡上优化 kernel 变得非常复杂。我大概负责了三个子项目,可以理解成三个算子,前前后后的代码实现加优化花了差不多一年的时间。 其中最大的困难是怎么设计完整的优化方案可以在一些边角 case 上也能够获得足够高的性能。
比如一开始的子项目是做一个 reduce 的算子,给一个最高五维的 tensor,里面有两个维度需要 reduce。这里面需要考虑的因素有:
一是最后一个维度是否需要 reduce,这涉及到怎么做向量化的问题。
二是在哪个维度做并行,如果是需要 reduce 的维度做并行的话,会有一次多个大核结果累加的问题,需要尽量避免。但是有的时候又无法避免,因为要考虑需要 reduce 的维度都很大,不需要 reduce 的维度都很小这种极端的特例。
三是在哪个维度开始分块的问题,比如 [x1, x2, x3, x4, x5]。在 x3 维度分块的话,比如切分出一个 32 的维度,意味着每次搬运都会搬 [32, x4, x5] 这么多的数据,这又得考虑这么多数据是不是在 cache放得下。不在这个维度做分块的话,又在哪个维度进行分块,怎么做是最优的。
诸如此类,需要考虑的极端情况很多。
QA 会随机生成各种 case,凡是达不到性能指标要求的都要写报告进行说明,最后明确是硬件本身的问题导致性能确实上不去,这个工作才算完。为了这个玩意我写了大约二十种条件分支来覆盖所有的边界情况。写完之后发现 code size 太大了,icache 总是 miss,产生一个抖动的现象,性能一下子又降到只有正常的 30%。又只能通过各种方式降 code size,然后来回折腾,干了两个月的时间,才全部测试 ok,通过相应的性能指标。
随后又做了一些其他的优化项目,整体来说,都比较困难。但在其中得到的经验让我对优化这个事情有了更加深刻的理解。觉得所有硬件的优化都是相通的。后面在看英伟达架构上的一系列变动和优化都能很快地领悟。 主要矛盾都在于访存能力跟不上算力导致的 gap,在计算部分能够打满算力的时候需要优化各级的数据搬运来获得更好的效率。 后续的工作也都是以此为指导。
在完成了 A 公司的优化任务之后,大概在 2023 年年中的时候, 我们开始做大模型推理的工作,自研了一个大模型推理框架。先是在端侧发力,在高通的手机上跑通了 llama 的 7B 模型,优化的还算不错,做了些 demo。然后开始做云端的芯片,主要是针对国产芯片进行优化。 这里面包括 B 公司和 C 公司两家的芯片。B 公司的芯片跟 AMD 早期的 GCN 架构一致。第二代卡没有矩阵计算单元,跟 GCN 架构可以说是一模一样。第三代卡自研了矩阵计算单元。针对 B 公司家的卡,我们主要优化了第二代卡上的性能。
在硬件架构上,B 公司的芯片还是类似于 SIMT 那一套东西。软件架构上采用的是类 CUDA 的编程方式,以往的 CUDA 经验可以无缝衔接过来。在将代码跑起来之后做了一些 profiling,发现 推理的主要瓶颈都集中在矩阵乘部分。如果上下文非常长的话,attention 会比较耗时。 于是针对这两个部分进行了针对性的优化。
关于矩阵乘的部分,有HGEMM 和 INT4GEMM,然后针对 prefill 和 decode 的不同阶段,矩阵乘优化的重点会呈现出一些差异。所以主要有 4 种组合,关于 HGEMM 的部分,可以直接调用 blas 库,但是我们测下来 B 公司的 blas 库性能较差,还有很大的优化空间。所以最后 4 种组合都要针对性优化。
针对 prefill 的情况,对于 HGEMM 的话,官方的blas 库其实对于 MNK 非常大的情况性能是 ok 的,但是对于 N 不到 8k 的情况,还是有一定的优化空间,尤其是 N 越小,blas 库的性能越差。为了提高这些 case 的性能,先自己糊了一个 HGEMM 的 kernel,MNK 维度都比较大的时候性能能跟官方的 blas 库差不多。然后分析出主要的性能瓶颈卡在 gmem 到 smem 的数据搬运上,采用了 block swizzle 也没有办法完全解决这个问题,于是对权重矩阵进行了重排,采用了一种访存更加高效的 layout 从而提高了这些 case 的性能,即使在 N 小一点的时候也能够获得 90% 以上的计算效率,比官方的 blas 库快了很多。至于 INT4GEMM,由于是计算密集型,所以一般做法是将权重矩阵转换成 FP16 或者 BF16 的数据类型,然后再调用对应的 blas 库。为了提高性能,我们也做了一些尝试,在 HGEMM 的基础上又糊出了一个 INT4GEMM,在有些 case 上还是能够获得比直接转换更好的性能。
针对 decode 的情况,GEMM 都是访存密集型的算子。HGEMM 的实现,尤其显存小,batch_size 非常有限, 所以我们采用了一种针对瘦高矩阵乘的优化。具体参考的是《 TSM2: Optimizing Tall-and-Skinny Matrix-Matrix Multiplication on GPUs》 这篇论文。 而至于 INT4GEMM,也是针对性地糊了一个 kernel,性能还不错,但是发现很多时候瓶颈都卡在反量化上,但是 B 卡上的指令集也没有类似于英伟达 LOP3 这样的指令,所以反量化的性能并不是太好。后续还有挺多能做的工作,但是由于时间限制,没有做更深入的优化了。
针对 attention 的算子,觉得 B 公司的 kernel 性能不太行,又开始糊算子。decode attention 的算子就糊了好几版实现,最后用的 trt decode attention算子做了一些适配。prefill 阶段的 attention,糊了一个 flash attention,调了调,性能还不错,但是这个卡算力实在太差,在很长的上下文下,模型还是跑得太慢了。
总而言之,由于 B 公司的编程方式非常趋近于 CUDA,所以写 kernel 还是比较流畅,当时糊了一大堆算子。从大模型推理的角度来说,B 公司卡的算力比较低,所以提高模型推理性能基本上就等于提高核心算子的性能。这跟 N 卡还有一些不同,N 卡在 decode 阶段的提速还得考虑到 cuda graph 这种东西的影响。
B 公司的工作完成地差不多之后,开始做 C 公司的工作。C 公司的这款硬件对标的是 N 卡 A100,性能挺好,具体的就不方便多说。硬件架构上采用的就是 SIMT+tensor core 的这种形式。软件栈上几乎是对标 CUDA,可以通过设置一些编译参数让 CUDA 的代码跑起来。
当时接到这个任务是在 10 月中旬,老板说已经跟人家公司聊了。给了一些模型范围和测试 case,说我们要赶紧开发一下,然后推理性能要比别人强。时间紧任务重。既然接到了任务,就只能和组里的小伙一起努力。Flash Attention 的实现直接调用人家写的官方库,Decode 阶段的 attention,我们有自己的一个版本。然后 GEMM 部分,甲方给了一个 C++ 的版本可以用来修改,但是折腾了一段时间后发现有的 case 性能上升,有的 case 性能下降,而且 C++ 版本性能本来就差于带汇编的版本。所以只能跟老板汇报这么点时间想要深度调优根本来不及。
大算子的优化在短时间内搞不定,于是又开始扣一些小算子,然后尽可能地再做一些融合。虽然提升有限,但还是有一定的性能提高。然后也发现了一些有意思的东西,比如这块卡上的 SIMT 调度机制似乎做得不是很完善。有的时候切了更多的并行,反而会导致性能的下降。然后开更少量的 CTA,让每个 CTA 负责更重的计算任务性能反而好一些。可能是 block 调度的开销比较大。当然这一结论没有进一步地去跟进了。
总之折腾了一个月,大概到 11 月中旬。我们测出来的性能结论是 prefill 速度要慢一些,大概能到人家 92% 的性能,decode 速度要快一些,平均大概有 37% 的性能提升。最后拿着结果跟人家去聊,然后甲方爸爸说他们又更新了一下 vllm 的版本,做了一些优化。得拿他们的新版本再去做性能对比。然后一测发现 prefill 和 decode 速度都比人家的慢。这个事情也无可奈何。总之我们的项目一直没啥营收,还得长期投入,也确实难搞。于是我又得干点能赚钱的项目,比如做 D 公司的算子开发和优化。
D 公司是一家我没有听说过的芯片公司,那个时候我才发现原来现在的国产芯片公司居然有那么多。再说过 D 这家公司,硬件架构上,他们不太愿意跟我们透露太多的东西,总之也可以简单地理解成几个大核拼在一起。 但是出现了一个比较难受的 layout 概念,这个玩意一般来说是搭配 tensor core,但是这个卡上所有的算子都需要跟这玩意打交道,一些简单的算子也的考虑相应的排布和对齐。我不太清楚这是硬件上的限制还是说软件上抽象了一些,就是为了更好地跑 SIMD? 不理解。软件架构上,玩的很花。但是主体是 MLIR,他们想通过这玩意来解决从 graph 到算子,再到底层指令的所有问题。但是水平不够,所以做出来的东西难以评价,非常抽象。
关于这套软件架构,从底向上来说吧。芯片有指令集,然后针对这些指令集,向上封装成一些类似于 Intrinsic 的玩意。计算指令好搞,但是访存指令,由于那套 layout 的影响,搞得十分乱,而且也没有抽象出足够合理的访存 API 供上层使用。所以后面在算子层,可以看到各自为战,大家都拿着底层的细粒度的数据搬运操作自己去凑出一堆更粗粒度的数据搬运操作,代码写得乱七八糟,令人无语。
再到算子层,为了解决算子中的并行和分块问题,他们想用一套更自动化的方式。可以简单地将算子理解成两个部分,第一个部分是 API 接口,用的是codeGen 的方式,针对不同的输入去动态地选择不同的并行方式和分块方式,然后用多少片上存储也是在这个时候确定。
第二个部分是 API 实现,可能是没太多编译器的人,kernel 实现只能用 C 语言,cpp 都没法用。kernel 里面拿到相应的并行和分块信息,然后完成相应的计算任务。这套玩意非常像 JIT,比较难受的是中间结果,所有的运算,函数接口都是传 Tensor 一样的玩意。如果碰上 page attention 的那种算子,指针传个两三次得到地址,然后再做数据搬运就会极其费劲。总而言之,他们想解决一些问题搞出来这套玩意,但是最后根本解决不了实际的问题,反而带来了开发复杂性的急剧提升。再往上走是 graph,框架相关的东西。graph 倒是正常地做各种 lower,反正功能是有的。然后往上是框架,pytorch 的适配,vllm 的适配,感觉做的都是一地鸡毛。这固然有时间的原因,很多东西短时间内做的不好是正常的。但是很大程度得归咎于软件架构设计的问题。前期拍脑子想方案,后期只能硬着头皮屎上做雕工。
再多说点,他们想拿这个卡做大模型推理。但是迟迟搞不定调度、graph、算子这些东西耦合在一起导致的问题。进度非常慢,即使 23 年的大模型调度策略都不能正常用起来。 里面的老哥工作地非常幸苦,无穷无尽的加班。 但是做这玩意的时间越长,发现离业界越远。这充分证明了团队没有足够的技术能力的时候,还是不要去玩花活。好好地做一套类 CUDA 的东西,再适配适配开源的东西。然后有点销售上的门路,日子还能过得不错。
一些思考
关于国产芯片和 NV。 先说一个核心问题,国产芯片有没有机会,或者说没有太强政企关系的国产芯片有没有机会。针对这个问题,我不是什么专家,只是在从业者的视角下说一下我的看法。有没有机会,当然是有机会的。 只要抓住了应用层,就会有成为真正中国英伟达的希望。 在应用角度,以大模型为例。用户接到了两个 token,他是认不出这两个 token 哪个是来自国产芯片,哪个是来自英伟达。在这种情况下,最重要的就是快。用户虽然认不出哪个 token 是来自国产芯片,哪个是来自英伟达。但是用户知道哪个是快的 token,哪个是慢的 token。 关于这个快,再引申一下,其实是有两个层面上。第一个是性能上的快,第二个是适配新技术的快。
关于性能上的快,这其实是一个多维度的东西,这涉及到模型本身的大小,硬件集群的机器数量等等因素影响。我们来控制一下变量,就单论单卡的性能,每张卡的算力以及对应的带宽。当然这个东西也受限于制造工艺的影响,成本的因素。 总的来说是要解决算力和功耗的问题,计算和存储的速度配比的问题。这个性能上的快,除了硬件上的问题,还有软件上的问题。硬件性能标称多少 T,但是软件上是否能发挥出来也是一个巨大的问题。 用户感受到的速度是硬件性能×软件发挥的水平。而基于 CUDA 生态,英伟达在软件上与其他家拉开的差距远比硬件更大。
关于适配新技术的快,主要是涉及到生态问题。在这方面上,英伟达永远是最快的。 因为算法人员总是在 N 卡上去实验和验证算法,做 infra 的人总是热衷于在 N 卡上扣细节做优化。这个原因主要在于 N 卡的稳定性以及开源社区能够有一大堆的资料文档参考借鉴。这些东西组成了一个牢不可破的护城河。而对于国产芯片,永远要慢,至于慢多少,则差别各异。这个问题很难解决。因为路是别人走出来的,你要去走这条路,那永远只能跟在别人身后。新技术一直是基于 N 卡,那么不管再怎么努力,国产芯片都是追不上的。再来回顾英伟达在多年前追赶英特尔的状态,在芯片上堆计算单元从而拿到更好的算力。但是依旧用的人少,所以到处送卡,让研究者用他们的卡去做研究。去抓住新技术的风口。 AlexNet 引爆了深度学习革命,用的是 N 卡。OpenAI 推出的 GPT 掀起了大模型的热潮,用的是 N 卡。下一波即将到来的新技术在哪里暂不清楚,但是从动态的视角去看世界,永远有下一波新技术的兴起。以目前的视角来看,国产芯片在新技术上,只能是适配,不断地去适配新的模型,新的框架,新的优化方法。什么时候能站在潮头,尚在虚无缥缈间。
关于国产芯片的硬件和软件。 前面将国产芯片做了一个泛化,然后和 NV 做了一些对比。但国产芯片也是一个很大的概念。 目前做国产芯片的公司大概有好几十家,知名度相对高一点的大概就十家左右。路线上主要是 SIMT + GEMM 单元这样的方式,摸着 NV 过河。 软件上类 CUDA 的编程方式,矩阵计算单元搞点汇编来调用。整个软件栈就长成 NV 的模样,大量的开源资源可以利用。额外做些适配即可。但是也有其他的方式,硬件上采用 SIMD+GEMM 单元的方式,然后软件上比较发散。 各式各样的编程方式都有,计算库、编译器、工具链再到上层的框架适配,所有的玩意都得自己搞一套。但是由于编程性比较差,很难去适配新技术。然后因为非主流技术,人员不断流失,进一步地影响新技术的引入。 写个算子再做做优化,要两三个月,写完之后发现大家又开始研究新的玩意了。恶性循环一旦启动,就难掩颓势。唯一能做的就是堆海量的研发人员加班加点尽可能地跟上,但是这需要非常大的成本,这也不是一般公司能干的事。
关于国产芯片公司能不能去 。其实国产芯片这个赛道,有很明显的机会,也有很明显的弊端。机会在于政治环境,国家需要。但是谁能真的成为中国的英伟达,目前来看,也不是很好说。但总体来说,目前我觉得后来者的机会不是很大了,但也会有一些市场,只要能有稳定的资金流和甲方爸爸,现在入场开始都来得及。总体来说,现在还是一个激烈角逐的市场,赢的人就有机会从英伟达几万亿的市值里面分到一杯羹,这是一个巨大的诱惑。说完机会,再说弊端,弊端主要是两个部分, 一个部分是由于出货量少,相应的直接招的岗位会少。就业市场上可能会有些劣势。很多公司招国产芯片的人一般是买了相关的卡,出货量多的自然对应的岗位会多一些。第二个弊端主要是一些芯片公司用的技术相对老旧,也比较封闭,跟不上新技术的发展,也比较困难。
关于个人职业发展。 这三年围绕着高性能计算或者说 AI infra 领域做了一些工作。除了上述所述的硬件,也搞了很多其他的东西。涉及到 NV、AMD的芯片,高通、Arm 的移动端 GPU、 Imagination 等等。深入的程度不一,但大部分时间都是疲于去完成各种项目。自己的技术水平似乎并没有太大的提高。而且在一些芯片上费老大劲做的东西,没人知晓,没有人用,更谈不上认可。还有一些工作,甚至从立项本身都没什么意义。所以这导致我经常处于焦虑的状态,比较长一段时间在工作中没法获得成就感,也看不到什么希望,就好像处在代码的泥沼里打滚。后续尽量做一些更有价值的工作,潜心多提高自己的技术水平。
来源:小镇评论家