ViT模型片外访存分析

对于ViT-tiny模型，其embedding大小为192。输入的patch_size为16，对于一张(256, 256)的图片而言，其总共包含1025个token(算上CLS token之后)。

考虑一个边缘端芯片，其总共的memory size一般<1Mb。若input_activation的memory与weight memory都为512Kb，且IA量化为8bit，weight量化为4bit的情况下。统计两个memory各自的访存。

ViT-tiny backbone layer

计算$Q,K,V$

将input feature map以170个token为1组放到片上，此时weight memory可以存下所有的QKV矩阵。生成的Q, K, V立即发出片外。总共的offchip access为 $ 170 \times 192 \times 8 \times 4 \times 6 = 11.95 Mb $，其中8为input activation (IA)的quantization bits, 4为load input feature map, offload Q, K, V这四次offchip access。6是1025个token可被分为6组。

计算$QK^T$

将n个Q与K token存入input SRAM，经过计算得到大小为$n^2$的矩阵，解方程 $ ( 2 \times 192 n + n^2 ) \times 4 = 512 \times 1024 $, 解得 $ n=128 $，故input SRAM EMA为$ (128 \times 192 \times 8 + 1025 \times 192 \times 8 + 1025 \times 128 \times 8) \times 8 = 21.52 Mb$, 其中括号外的8意为1025个token分为8组。

计算$PV^T$

将n个P与V token存入input SRAM，经过计算得到大小为$n \times 192 $的矩阵，解方程 $ ( 2 \times 192 n\times + n^2 ) \times 4 = 512 \times 1024 $, 同样 $ n=128 $，故input SRAM EMA为$ (128 \times 128 \times 8 + 1025 \times 192 \times 8 \times 2) \times 8 = 25.02 Mb$

计算$FCo$

将input feature map以113个token为1组放到片上。生成的Output Feature map与上一个layer的output feature map做short cut后立即发出片外。总共的offchip access为 $ 113 \times 192 \times 8 \times 3 \times 10 = 4.96 Mb $，其中3为load input feature map, output feature map与short cut feature map这三次offchip access。

计算$FC1$

将input feature map以68个token为1组放到片上。生成的Output Feature map大小为[68, 768]。由于weight matrix大小为[192, 768]，以4bit量化为例，存下所有的weight matrix需要576Kb的片上memory。考虑用一些压缩技术使片上memory可以存下整个weight matrix，则总共的offchip access为 $ 68 \times 192 \times 8 \times 5 \times 16 = 7.97 Mb $，其中5为input feature map与其4倍大的output feature map的offchip access。

计算$FC2$

将大小为[48, 768]的input feature map以48个token为1组放到片上。生成的Output Feature map大小为[48, 192]。总共的offchip access为 $ 48 \times 192 \times 8 \times 21 \times 7 = 10.34 Mb $

将以上的EMA加总，得出计算一层ViT-tiny backbone layer input SRAM的EMA为54.27Mb。相比之下，在每一次inference weight只加载一次的情况下，weight EMA为$192 \times 192 \times 3 + 192 \times 192 + 192 \times 192 \times 4 \times 2 = 5.06 Mb $
可见由于ViT input token length较大，带来了很多input feature map反复读取的开销。
取典型的加速器throughput为200GOPS，整个backbone layer的MAC数约为1.526 GOP，则不考虑EMA的情况下，计算一个backbone layer的latency约为7.63ms。
取10pJ/bit为片外access 能耗，若要在10ms的区间内搬运54.27Mbit，平均比特率约为5.43Gbs，计算得功耗约为54.27mW。