Transformer~45

这是Transformer时代端侧AI芯片的机遇与挑战  , 还是先感谢大佬啊 搬来给自己学习 勿怪啊

Transformer似乎已经将CNN的地位逐步吞噬，而Transformer需要更大的计算量，更容易被带宽和缓存掣肘，那么短策AI芯片公司该如何应对这一机遇和挑战？本文分析总结了各个模态典型transformer模型及其在端侧视角下的优劣。

似乎就在这两年，transformer在语音、视觉、雷达这些模态对CNN悄无声息就形成了碾压之势。考虑到这些领域的模型几乎承包了端侧的workload，那么，广大端侧AI芯片公司接下来该如何应对？下一代芯片设计选型该如何考量？我想业界的朋友们肯定没少调研：Transformer万般好，但在端侧这个算力规模下，相比CNN，Transformer往往需要更大的计算量才能博得实质性的优势，同时Transformer的计算强度更低，更易被带宽和缓存掣肘。。。所以需要采用和面向LLM的大算力AI芯片类似的策略：加缓存，加带宽。

额，这样真的就够了？两个月前的我还真是这么觉得的，感叹transformer大一统的时代已来，堆料就好，甚至我在下面这个回答中还输出了些许对AiSys这个领域接下来发展的悲观态度。但最近，在研究了这些领域大大小小热门transformer结构，以及和一些算法的朋友交流后，我发现transformer在其他模态上的应用范式和NLP存在不小的差别，妄图将面向LLM的芯片设计决策直接套用到端侧并不合适。 

接下来我将罗列各个模态典型transformer模型及其在端侧视角下的优劣，思考与总结将在最后给出。对模型不感兴趣的朋友可以跳过铺垫，直接看总结。

此外，为了更好的交流（和谐），一如既往先对讨论对象下定义，本文所指的端侧AI芯片代表下图这个圈。Mobile、Edge、Autonomous。这个区间的AI芯片，一方面对成本更加敏感，在缓存和带宽上的堆料相对保守，另一方面能效更有优势的DSA占比更高，导致相对编程灵活性偏弱。端侧典型transformer模型盘点

视觉主干篇

Vit： AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE, 2020

简介

• 整体流程上：通过patch embedding将图像直接打包编码，转换成transformer类型的序列特征，接上patch position encoding后，后续由经典的transformer block全权接管。

• 算法角度上：Vit结构只通过少量特征前处理就把视觉信息转换为了序列特征，还是挺优雅简洁的。但由于尺度信息在patch embedding阶段就被patch_size唯一确定，致使对视觉任务至关重要的多尺度信息缺位。因此，在除分类之外的其他视觉任务中，Vit的表现并不理想。

端侧视角下的优劣

• 优势：

  • Vit结构除了patch embedding与输出形式之外，与经典transformer结构并无二致，所以可直接享用抛开kv-cache的绝大多数transformer推理优化策略。

• 劣势：

  • 计算量随分辨率二次增长(seqlen=(H/patch_size)x(W/patch_size))。当应用到大分辨率场景时（比如检测分割），无论是FFN还是MHA中的矩阵乘都无法直接载入资源受限的端侧近核缓存。虽然有各类并行切分策略支撑，但缓存与主存间的超额换入换出产生的带宽开销还是严重制约了它的计算效率。

  • patch_size级别(16/32)的特殊conv2d，也并不一并有很好的利用效率。当然重点还是1)，2)的问题和CNN第一个inchannel=3的卷积很像，不是瓶颈，而且2)也可以转换为permute+gemm的形式规避。

swintransformer：Swin Transformer: Hierarchical Vision Transformer using Shifted Windows, ICCV, 2021简介

• 整体流程上：参考上图，PatchMerging逐Stage降尺度(第一个PatchPartition+LinearEmbedding可以看做特化的PatchMerging)；单个SwinTransformer先通过WindowPartition将视觉特征转换为以局部窗口为粒度的序列特征，送入经典Transformer结构处理完毕后，再通过WindowUnpartition转回视觉特征；每俩SwinTransformer做一次窗口滑动，以实现全局的信息流通。

• 算法角度上：SwinTransformer这套魔改的WindowAttention还是肯定了视觉任务中的局部性特点。另外，在整个pipeline中视觉特征与序列特征的来回切换，使得基于注意力的多尺度信息提取成为可能。这对Transformer在检测、分割等视觉任务上的落地有重要意义。

端侧视角下的优劣

• 优势

  • WindowAttention机制使得计算量随分辨率线性增长，解决了ViT的主要问题。此外，序列规模seqlen仅与窗口尺寸ws相关（seqlen=ws x ws），只要embed_dim设计的当，无论何种分辨率，都能将矩阵运算常驻于AI芯片的近核缓存。

  • 该机制还变相地提升了Batch的大小(Batch=batch_origin x (H/ws) x (W/ws), ws指窗口大小)。并行无需同步，完全免费！！这对有多核的架构设计来说，绝对是意外之喜。

• 劣势

  • 除了继承了ViT的PatchEmbedding的特殊卷积之外，WindowPartition/Unpartition和窗口滑动引入了permute/torch.roll，和windowsize与分辨率未匹配引入了Pad，这些高频的数据重排无疑对性能有较大的影响。

后续的演进

• 注意力机制上：窗口注意力（SwinTransformer）与全局注意力（ViT）的结合。比如ViTDet：Exploring Plain Vision Transformer Backbones for Object Detection, ECCV, 2022。就是部分Transformer使用WindowAttention，部分使用Attention。

• 位置编码上：将NLP领域的位置编码创新往视觉任务上套，比如RoPE: ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING, 2021。如下所示，核心计算上就是奇偶旋转偶去反，(x, y) --> (-y, x)。在onnx层级会被拆分为：低维split+neg+低维concat的子图。将一个明显的1pass转换为了multipass，对性能还是略有影响的。

def rotate_half(x):
    x = rearrange(x, '... (d r) -> ... d r', r = 2)
    x1, x2 = x.unbind(dim = -1)
    x = torch.stack((-x2, x1), dim = -1)
    return rearrange(x, '... d r -> ... (d r)')

视觉检测篇

DETR: End-to-End Object Detection with Transformers, ECCV, 2020简介

• 整体流程上：如上图。主干提特征，做一道ChannelMapping和位置编码后，送入经典的Transformer编解码模块，并最终交由FFN出框。值得说明的是，解码部分的query embedding是一个预设的固定值（学习参数）。

• 算法角度上：基于transformer的端到端框架，无NMS，无任何人工特征与先验知识，但训练困难。部署视角下，简直是惊为天人的设计。

端侧视角下的优劣

• 优势

  • 端到端没有nms，后处理无需offload到cpu。

  • 和ViT一样，基本上能无缝使用绝大多数transformer推理优化策略。

• 劣势

  • 从检测的典型负载而言，这个计算量还是偏大。达到同样的算法效果，相比CNN时代检测头组合，开销太大，还是有点得不偿失。

Deformable DETR: DEFORMABLE DETR: DEFORMABLE TRANSFORMERS FOR END-TO-END OBJECT DETECTION, 2020简介

• 整体流程上：Deformable DETR的流程有点复杂，我是扫了代码转了模型，才品出门道。感觉单凭上面两张结构图，似乎不容易搞懂。我尝试以模型结构的视角简练地解读一把：Deformable DETR接收主干提取的多尺度视觉特征，送入其定制的DeformableTransformer结构中。在该结构内，query特征存在视觉多尺度和序列两种表示形式。基于序列特征预测感兴趣特征点的坐标（Fig2中的referencepoints）与偏移（Fig2中的offset），再按该坐标分别从多尺度视觉特征中抽取特征点并结合序列特征计算得到的权重，得到下一轮DeformableTransformer的query输入。

• 算法角度上：降低了训练的难度。

端侧视角下的优劣

• 优势

  • 好像很难给Deformable DETR找一个部署层面的优点，硬要说的话，其魔改的DeformableAttention矩阵计算确实要少一点。

• 劣势

  • 如果WindowPartition和WindowUnpartition在模型中不断插入的数据reorder(permute)对端侧性能有冲击的话，我想DeformableAttention的高频grid_sample抽取特征点的操作，对不灵活低带宽的端侧只能用灾难来形容。

  • 更不要说多尺度视觉特征与序列特征切换时引入的特化split（非等分切分）和高维permute了。

后续的演进

• DINO: DETR with Improved DeNoising Anchor Boxes for End-to-End Object Detection, 2022，在DeformableDETR基础上进一步引入了query selection，在解码阶段增加了目标优选（topk+gather）。

点云篇

DSVT: DSVT: Dynamic Sparse Voxel Transformer with Rotated Sets CVPR 2023简介

• 整体流程上：整个DSVT的pipeline（上图中的DSVT部分）可以理解为点云版本的swintransformer。1). 可并行的局部自注意力：DSVT将稀疏voxel分为多窗口并最终分为等数量的set，针对每个set做attention。2). 全局信息流通(shift窗口滑动)：DSVT分别对x/y做partition。

• 数据的切换方式：预先记录在x/y方向排序对应至voxel特征上的索引，并通过gather/scatter实现数据在voxel点云特征空间和transformer序列特征空间之间的切换。

• 流程中的Voxelization和BEV都是领域专业术语，本文不做讨论，voxel部分可参考VoxelNet: End-to-End Learning for Point Cloud Based 3D Object Detection CVPR 2018。

• DSVT和SwinTransformer除了流程之外，典型负载也比较相似（DSVT中set）

端侧视角下的优劣

• 优势

  • 继承了WindowAttention的所有收益。降低了计算量、单一矩阵的计算规模，提升了计算的并行度。

• 劣势

  • 数据在transformer序列特征空间和voxel点云特征空间下的切换成本远高于其与视觉特征的。一是scatter/gather这类数据随机访问(shuffle)远比数据重排操作(reorder)耗时，二是这些操作操作在端侧的支持性是个问号。全部offload给cpu吗，这不是开玩笑。

  • 一句话，端侧不友好跟DSVT本身没关系，是基因不好吃了点云数据结构的亏。

3D检测头的演进

3D检测头的算法演进路线比较复杂，要做的铺垫很多，比如CNN时代的CenterNet、CenterPoint，但结局却与2D检测殊途同归，就是DETR化。所以反而不用着墨太多。

多模态篇

CLIP: Learning Transferable Visual Models From Natural Language Supervision ICML 2021简介

• 整体流程上：图片编码和文本编码部分各经过transformer encoder（也可以是其他编码结构），再对各自的编码特征做多模态融合（norm），再计算彼此的余弦相似度（matmul, cos）。

端侧视角下的优劣

• 优点

  • 和ViT一样，简单质朴的结构，大道至简的践行者，这对端侧落地是一个利好。

• 缺点

  • 包含prompt engineering的text encoder，其输入文本序列会随应用上下文实时改变，对端侧软件栈的动态shape支持性提出了更高的要求。

轻量级Transformer

从模型结构角度来说，我把轻量级Transformer分为两个派系：有机结合派与深入改革派。

有机结合派

该派别的代表的有：

• EfficientFormer: Vision Transformers at MobileNet Speed NIPS 2022

• MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformerarxiv 2021

只需从Architecture Overview中我们就能发现他们的端倪：熟悉的conv+relu+bn，逐级降采样。这类模型还是保留了CNN架构的典型特点，只有头部和躯干部分植入了Transformer单元。

这类算法虽然可以灵活地切换血统（调整其Transformer单元的占比），来达到在算法效果和推理性能上的平衡。但是从系统的角度出发，我感觉这类算法在定位上比较模糊：究竟是面向CNN算力区间，Transformer算力区间，还是两者的中间地带？至少到目前为止，从我了解的信息来看，中间地带是这个派系的主要舞台。

深入改革派

该派别的代表的有：

• Next-ViT: Next Generation Vision Transformer for Efficient Deployment in Realistic Industrial Scenarios, 2022

• EfficientViT: Lightweight Multi-Scale Attention for On-Device Semantic Segmentation, ECCV, 2022

这个派系已经不满足于Transformer粗浅地植入，他们更深层次地改造Transformer结构，用pooling裁剪注意力数量以缩减注意力矩阵规模，用relu/pooling实现近似softmax的全局注意力机制。

相比于有机结合派，这个派系迁就现有AI芯片特性的算法改造，反而目标清晰，就是要在CNN的算力区间里一决雌雄。设计“Domain Specific Algorithm”，对于部署和优化来说，自然是再好不过了。但是这种算法向平台的反向适配，带着旧时代AI芯片枷锁负重前行的想法，我总感觉与主流的创新思路有点背道而驰。

总结

更好的编程性

归纳偏执的差异深刻的影响了CNN和Transformer的演进模式，CNN的局部性特点无疑更多地引导了网络全局结构层面的创新，比如ResNet、FPN等等，而Transformer的全局性则反过来使改进优化更多体现在了针对Attention和FFN细节上的修修补补，也就是算子层面。再叠加上Transformer在模型中往往以Block/Layer形式堆叠，那么围绕Transformer所引入的算子创新将复制到所有Transformer单元。因此，Transformer时代，这些模态模型特化算子的占比将会远高于过去。上述盘点的DeformableDETR/DSVT的grid_sample/gather/scatter与高维张量，SwinT和VitDet引入的WindowPartition与Unpartition操作（额外的Pad与高维Permute），以及RoPE等特殊的位置编码运算，就是最好的例证。

这个时候，倘若仍按照CNN时代的思路，将grid_sample/gather/scatter等编译软件栈无法处理的计算负载offload到cpu，将高频地打断pipeline并额外挤占带宽，必然大错特错。相反，配置与TensorCore共享高速缓存的可编程加速单元，或者在DSP中嵌入TensorCore都会是一个更好的设计方向。

考虑到应用FlashAttention这种对计算灵活性要求极高（我这里说的是量化版FA）的优化策略，上述决策只会被进一步强化。那么，进一步联想，在这个趋势下，走端侧路线的存算AI芯片是否会面临更大的挑战（毕竟不是所有运算和优化都适合量化）？whaosoft aiot http://143ai.com

更合理的资源配置

更贴合的缓存尺寸。无论是受限于WindowAttention的视觉主干和点云DSVT，还是受限于Detection Objects的检测任务QueryEmbedding，亦或是多模态精简的prompt，这些算法在设计上都暗示了序列长度并不会随模型计算量增大而无限膨胀。简单算一下各模态transformer较大的参数配置场景（mlp_ratio=4, queryobj=seqlen=1024, embed_dim=1024），8M高速缓存就能把完全兜住单batch8bit推理时任何单一计算的权重数据。所以没必要无限堆料，尺寸到位就行。

有潜在优势的多核设计。从SwinTransformer/ViTDet等workload来看，在强调局部性的任务中算法建模有更高的意愿将Transformer的注意力输入改造代表局部特征的相互独立的子序列。在这种情况下，多核设计的收益将更加的明显。

更精准的算力定位。 目前来看，transformer时代算法存在两种创新模式，一种是轻量化Transformer，与CNN深度结合，在原有结构不变的前提下，引入少量Transformer模块，走底层路线；另一种是基于模态特点，引入新算子适配特征数据，做Transformer定制化，走的是上层路线。路线之争我不敢妄议，但有一点可以肯定，这两者的算力区间是截然不同的。一款AI芯片的设计定位，不能摇摆模糊，只能二选一。两头都要，底层路线没有价格优势，上层路线算力资源又跟不上，结果可能会比较难看。因此，需要团队里的软硬件协同研究员们在设计阶段的更多投入，把握好transformer效果优势的边界点。

更精细的Buffer建模

这些workload充分体现了：

• 非语言模型跨特征空间的求解述求。数据会在序列特征和其他模态特征上来回切换，以寻求更好的特征抽取。比如SwinT中图像特征下的shift与降采样，以及序列特征下的Attention与FFN，再比如DSVT在点云特征下的特征分组，和同样序列特征下的Attention/FFN。

• 位置信息的花式编码述求。比如RoPE，以及DINO中的这个解码端position操作。

若IR层级的Buffer建模能从tensor颗粒度直线跨越到item颗粒度，那么一方面，跨特征空间切换所引入的reorder(permute)和shuffle(grid_sample/gather)操作将能被更好的刻画，从而拥有了进一步挖掘其与前继与后继的eltwise/gemm节点融合的空间，另一方面，花式位置编码(RoPE等)计算子图中各Buffer间的集合和依赖关系也将被更好的建模，从而使整个子图从multipass转换为1pass成为可能。这些性能收益的叠加很可能成为最终性能对标时的胜负手。

更高扩展性的软件栈

Transformer算法的这些特性对软件栈的支持性上也有一定的影响：

• 各模态Transformer引入的特殊算子中，除了shuffle类操作，绝大多数workload(高维张量运算，高维Permute，非等分Split等)都是编译软件栈的权责范畴，由它将这些负载平凡化并完成硬件适配。

• 虽然除了多模态之外的workload并没有对动态Shape存在强烈诉求，但鉴于多模态模型所展示的巨大潜力，编译软件栈还是有必要提前布局动态Shape，并做好合理的版本规划。

而这些在CNN时代可能压根没被纳入测试用例，甚至是功能清单。显然，这是AI芯片编译软件栈的可扩展性提出了更高的要求。设计上的考量维度太多，但测试这个方面，我倒是觉得到了上马Fuzz的阶段了。

也许能从后续的SDK版本更迭和release note中窥探一下各家的应对情况