qq_49588762的博客

然而，目前广泛使用的加速器是英伟达的通用加速器GPU。GPU是以计算为中心的架构，使用SIMT掩盖访存开销。对于GEMM算子，GPU打不过TPU；对于GEMV算子，GPU打不过AiM。但是，GPU有很强的通用性，加上完善的生态，它可以很好地完成各种算子（不论是计算密集性算子还是访存密集型算子）。从加速器通用性角度看，AiM与TPU类似，两者都是十分专用的加速器。TPU使用脉冲阵列加速计算密集性算子GEMM；AiM使用存算架构加速访存密集型算子GEMV。是一个专门用于加速GEMV算子的加速器。

目前，ALU算的很快，把内存中数据加载到cache的带宽有限，这时候，ALU的算力就会被浪费，性能瓶颈卡会在内存带宽上，这就是“内存墙”。如下图，简单来讲，当thread 8运行的时候，其他的线程（比如thread 1, 3）也没闲着，它们正在取数据。然而，TPU是一个专用的加速器，只能用来加速GEMM算子，其他的活统统干不了。简单来讲，GPU的架构就是堆很多的寄存器和ALU，获取高算力。MLPs (Multi-Layer Perceptions) ，多层感知机，大家应该都挺熟悉的，里面的算子是GEMV。

如果sequence中context比较长，cuda block承担的任务量将会比较大，甚至有可能出现硬件资源（共享内存，寄存器）不够的情况。我们来看一下paged_attention_v1_kernel中的并行计算任务划分，其实就是搞清楚cuda block、warp和thread各自承担了那些计算任务。因此，paged_attention_v1_kernel将一个head中的attention任务交给一个cuda block完成。这样的话，不涉及不同cuda block间的同步，通信开销会比较小。

大语言模型（LLM）的高吞吐量服务需要一次批处理足够多的请求。然而，现有的推理系统面临着巨大的挑战。一方面，每个请求的键值缓存（KV cache）占用大量的内存，而且会动态增长和收缩，显存容量不足限制了批处理的大小；另一方面，attention算法中大量使用访存密集型的GEMV算子，显存带宽成为了系统瓶颈。为了解决这个问题，我们提出了PIM-attention。PIM提供了更大的内存容量和bank级并行的访存带宽，为加速attention算子提供了机遇。

在decode过程，推理出一个token的时延 = bank内并行访存（GEMV）的时延 + Controler-bank 通信的时延 + Controler 内访存（softmax、Norm 和 reduction 等）的时延。定量分析decode过程，也就可以分析出在给定访存带宽下，模型推理的速度。bank内并行访存（GEMV）的时延包含两部分：模型参数相关的GEMV的时延和kv cache相关的GEMV的时延。在decode阶段，Controller和bank内的算力均可以吃满访存带宽。