国产AI芯片三年，一些思考

公众号记得加星标，第一时间看推送不会错过。

来源：内容转自知乎，作者：中国科学院大学 计算机系统结构硕士 有了琦琦的棍子，谢谢。

刚好最近有时间，聊聊这三年在国产芯片上的工作经历。一方面是记录一下这三年做的一些事，一方面是基于这三年的工作经验引发的一些思考。我是在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的东西，再适配适配开源的东西。然后有点销售上的门路，日子还能过得不错。

一些思考

1. 关于国产芯片和NV。先说一个核心问题，国产芯片有没有机会，或者说没有太强政企关系的国产芯片有没有机会。针对这个问题，我不是什么专家，只是在从业者的视角下说一下我的看法。有没有机会，当然是有机会的。只要抓住了应用层，就会有成为真正中国英伟达的希望。在应用角度，以大模型为例。用户接到了两个token，他是认不出这两个token哪个是来自国产芯片，哪个是来自英伟达。在这种情况下，最重要的就是快。用户虽然认不出哪个token是来自国产芯片，哪个是来自英伟达。但是用户知道哪个是快的token，哪个是慢的token。关于这个快，再引申一下，其实是有两个层面上。第一个是性能上的快，第二个是适配新技术的快。

关于性能上的快，这其实是一个多维度的东西，这涉及到模型本身的大小，硬件集群的机器数量等等因素影响。我们来控制一下变量，就单论单卡的性能，每张卡的算力以及对应的带宽。当然这个东西也受限于制造工艺的影响，成本的因素。总的来说是要解决算力和功耗的问题，计算和存储的速度配比的问题。这个性能上的快，除了硬件上的问题，还有软件上的问题。硬件性能标称多少T，但是软件上是否能发挥出来也是一个巨大的问题。用户感受到的速度是硬件性能x软件发挥的水平。而基于CUDA生态，英伟达在软件上与其他家拉开的差距远比硬件更大。

关于适配新技术的快，主要是涉及到生态问题。在这方面上，英伟达永远是最快的。因为算法人员总是在N卡上去实验和验证算法，做infra的人总是热衷于在N卡上扣细节做优化。这个原因主要在于N卡的稳定性以及开源社区能够有一大堆的资料文档参考借鉴。这些东西组成了一个牢不可破的护城河。而对于国产芯片，永远要慢，至于慢多少，则差别各异。这个问题很难解决。因为路是别人走出来的，你要去走这条路，那永远只能跟在别人身后。新技术一直是基于N卡，那么不管再怎么努力，国产芯片都是追不上的。再来回顾英伟达在多年前追赶英特尔的状态，在芯片上堆计算单元从而拿到更好的算力。但是依旧用的人少，所以到处送卡，让研究者用他们的卡去做研究。去抓住新技术的风口。AlexNet引爆了深度学习革命，用的是N卡。OpenAI推出的GPT掀起了大模型的热潮，用的是N卡。下一波即将到来的新技术在哪里暂不清楚，但是从动态的视角去看世界，永远有下一波新技术的兴起。以目前的视角来看，国产芯片在新技术上，只能是适配，不断地去适配新的模型，新的框架，新的优化方法。什么时候能站在潮头，尚在虚无缥缈间。

2. 关于国产芯片的硬件和软件。前面将国产芯片做了一个泛化，然后和NV做了一些对比。但国产芯片也是一个很大的概念。目前做国产芯片的公司大概有好几十家，知名度相对高一点的大概就十家左右。路线上主要是SIMT + GEMM单元这样的方式，摸着NV过河。软件上类CUDA的编程方式，矩阵计算单元搞点汇编来调用。整个软件栈就长成NV的模样，大量的开源资源可以利用。额外做些适配即可。但是也有其他的方式，硬件上采用SIMD+GEMM单元的方式，然后软件上比较发散。各式各样的编程方式都有，计算库、编译器、工具链再到上层的框架适配，所有的玩意都得自己搞一套。但是由于编程性比较差，很难去适配新技术。然后因为非主流技术，人员不断流失，进一步地影响新技术的引入。写个算子再做做优化，要两三个月，写完之后发现大家又开始研究新的玩意了。恶性循环一旦启动，就难掩颓势。唯一能做的就是堆海量的研发人员加班加点尽可能地跟上，但是这需要非常大的成本，这也不是一般公司能干的事。

3. 关于国产芯片公司能不能去。其实国产芯片这个赛道，有很明显的机会，也有很明显的弊端。机会在于政治环境，国家需要。但是谁能真的成为中国的英伟达，目前来看，也不是很好说。但总体来说，目前我觉得后来者的机会不是很大了，但也会有一些市场，只要能有稳定的资金流和甲方爸爸，现在入场开始都来得及。总体来说，现在还是一个激烈角逐的市场，赢的人就有机会从英伟达几万亿的市值里面分到一杯羹，这是一个巨大的诱惑。说完机会，再说弊端，弊端主要是两个部分，一个部分是由于出货量少，相应的直接招的岗位会少。就业市场上可能会有些劣势。很多公司招国产芯片的人一般是买了相关的卡，出货量多的自然对应的岗位会多一些。第二个弊端主要是一些芯片公司用的技术相对老旧，也比较封闭，跟不上新技术的发展，也比较困难。

4. 关于个人职业发展。这三年围绕着高性能计算或者说AI infra领域做了一些工作。除了上述所述的硬件，也搞了很多其他的东西。涉及到NV、AMD的芯片，高通、Arm的移动端GPU、Imagination等等。深入的程度不一，但大部分时间都是疲于去完成各种项目。自己的技术水平似乎并没有太大的提高。而且在一些芯片上费老大劲做的东西，没人知晓，没有人用，更谈不上认可。还有一些工作，甚至从立项本身都没什么意义。所以这导致我经常处于焦虑的状态，比较长一段时间在工作中没法获得成就感，也看不到什么希望，就好像处在代码的泥沼里打滚。后续尽量做一些更有价值的工作，潜心多提高自己的技术水平。

*免责声明：本文由作者原创。文章内容系作者个人观点，半导体行业观察转载仅为了传达一种不同的观点，不代表半导体行业观察对该观点赞同或支持，如果有任何异议，欢迎联系半导体行业观察。

今天是《半导体行业观察》为您分享的第4117期内容，欢迎关注。

加星标第一时间看推送，小号防走丢

求推荐