1. 存内计算技术解析：从原理到TPU架构革新
存内计算（Compute-in-Memory, CIM）正在重塑AI加速器的设计范式。传统冯·诺依曼架构中，数据需要在存储单元和计算单元之间频繁搬运，这种"数据搬运瓶颈"消耗了高达90%的系统能耗。CIM技术的突破性在于将计算操作直接嵌入存储阵列，实现了"数据在哪，计算就在哪"的范式转变。
数字SRAM-based CIM的核心架构包含三个关键设计：首先是采用分bank组织方式，每个bank对应一个输出通道，实现并行计算；其次是在bank内部进一步划分sub-array处理不同输入通道，提升数据吞吐量；最重要的是在bitcell阵列中集成本地计算电路，支持原位乘加运算。以论文中提到的7nm CIM宏为例，其采用8T SRAM单元配合动态逻辑电路，在保持28nm工艺下实现了351TOPS/W@INT4的能效表现。
关键提示：CIM设计需要权衡计算密度与灵活性。纯模拟CIM虽然能效更高（如ReRAM方案可达78.4TOPS/W），但受限于工艺偏差和低精度问题。数字CIM通过保持数字信号处理的特性，更适合高精度生成式模型的需求。
2. 生成式AI模型的硬件挑战与计算特性
现代生成式AI模型呈现出两大主流架构：基于Transformer的大语言模型(LLMs)和扩散模型(DMs)。它们的计算模式对硬件提出了独特要求：
LLMs的双阶段特性：
- Prefilling阶段：处理整个输入序列（如1024个token），计算密集型，90%时间消耗在QKV生成、投影和FFN层的大矩阵乘法
- Decoding阶段：逐个token生成，内存访问密集型，Attention层占比提升至33.7%，GEMV操作成为瓶颈
扩散模型的变革： 传统U-Net架构正被DiT（Diffusion Transformer）取代，如Stable Diffusion 3已全面转向Transformer架构。DiT-XL/2模型的评估显示，99.31%的计算集中在DiT块中的Transformer层，其中Softmax操作意外地占据了36.9%的延迟，这源于高分辨率图像处理带来的超大attention矩阵。
实测数据揭示关键发现：在512×512图像处理中，DiT的GEMM操作仅占35.65%延迟，远低于LLMs的84.9%。这种计算特性的差异直接影响了CIM架构的设计取舍。
3. CIM-TPU协同设计方法论
3.1 基准架构建模
研究以TPUv4i为基线，构建了完整的CIM-TPU仿真框架。核心创新点在于用CIM-MXU替代传统脉动阵列MXU，同时保留其他组件（VPU、内存 hierarchy等）。关键参数配置体现设计权衡：
| 组件 | TPUv4i参数 | CIM-TPU参数 |
|---|---|---|
| MXU类型 | 128x128 脉动阵列 | 16x8 CIM核心阵列 |
| 计算精度 | BF16/INT8 | 支持FP/INT混合精度 |
| 内存体系 | 16MB VMEM + 128MB CMEM | 保持相同配置 |
| 互联带宽 | 614GB/s HBM | 维持相同带宽 |
3.2 CIM-MXU微架构创新
CIM-MXU面临的核心挑战是如何将小型CIM宏（通常仅128x256维度）组织成高性能计算单元。论文提出的解决方案包含三大关键技术：
- 二维脉动数据流：在16x8的CIM核心阵列中，输入向量沿水平方向脉动传递，权重矩阵沿垂直方向更新，形成输出固定的数据流
- 双模式精度支持：
- FP模式：将权重mantissa存入CIM阵列，配合预处理单元完成指数对齐
- INT模式：直通计算，利用CIM原生高效处理低精度运算
- 并发权重更新：通过专用Weight I/O端口，实现计算与权重加载的并行化，解决Transformer层权重复用率低的问题
在22nm工艺下的对比测试显示，该设计在保持相同16384 MACs/cycle吞吐量下，能效提升至7.26TOPS/W（9.43倍改进），面积效率达1.31TOPS/mm²（2.02倍提升）。
4. 面向生成式AI的架构优化实践
4.1 LLM专用优化（Design A）
针对LLMs的内存受限特性，采用4个8x8 CIM-MXU的保守配置，通过三项关键优化：
- GEMV加速：利用CIM的广播式计算特性，消除脉动阵列的流水线填充开销，使解码阶段延迟降低29.9%
- 动态精度切换：在Prefilling阶段使用BF16精度，Decoding阶段自动切换至INT8
- 稀疏支持：通过激活掩码跳过无效计算，实测在Alpaca数据集上带来18%的能效提升
4.2 DiT专用优化（Design B）
针对DiTs的计算密集型特点，采用8个16x8 CIM-MXU的激进配置，创新性地实现：
- Softmax硬件加速：在VPU中集成基于tanh近件的快速Softmax单元，延迟降低30.3%
- 条件注入优化：为shift and scale操作设计专用数据通路，减少53%的CMEM访问
- 分块Attention：将大尺寸attention矩阵分解为可拟合CIM-MXU的块，通过重叠计算隐藏数据搬运延迟
多芯片扩展测试表明，4个Design B TPU组成的环状拓扑，在512×512图像生成任务中实现33%的吞吐提升，同时MXU能耗下降6.34倍。
5. 工程实现中的挑战与解决方案
在实际芯片设计中，我们遇到了几个关键挑战：
信号完整性问题： 高密度CIM阵列导致电源噪声增加，通过以下措施解决：
- 采用分布式去耦电容布局，使IR-drop降低42%
- 开发自适应时钟门控方案，空闲bank自动进入低功耗模式
- 实现位线电压补偿电路，将计算误差控制在0.3%以内
热管理创新： CIM的高计算密度引发局部热点，我们的应对策略：
- 热感知任务映射：运行时监控温度，动态调整任务分配
- 异构冷却方案：3D封装中集成微流体通道，热点区域冷却效率提升3倍
- 温度导向电压调节：建立Vdd-Temp查找表，确保计算稳定性
工具链适配： 传统编译器无法有效映射CIM架构，我们开发了：
- CIM-aware的图优化pass，自动识别适合CIM计算的算子模式
- 双缓冲管理引擎，实现计算与数据搬运的100%重叠
- 精度损失分析工具，指导混合精度策略制定
实测显示，完整工具链可将硬件利用率提升至78%，远超基线TPU的53%。
6. 未来演进方向
基于当前研究成果，我们认为CIM-TPU架构还有以下发展空间：
异构计算集成：
- 探索CIM与近内存计算(NMC)的协同，处理不规则计算模式
- 研究可重构数据通路，动态适配LLM和DiT的不同需求
- 开发3D堆叠架构，突破内存带宽限制
算法-硬件协同：
- 设计CIM友好的模型架构，如基于块稀疏的Attention机制
- 开发自适应量化策略，平衡精度和能效
- 研究训练-推理一致的CIM架构，支持端到端优化
从产业视角看，CIM技术要走向大规模部署，还需要解决：
- 标准化接口定义，确保不同厂商CIM模块的互操作性
- 可靠性增强机制，特别是针对边缘设备的容错设计
- 跨平台编程抽象，降低开发者使用门槛
我们在后续工作中发现，采用CIM技术的TPUv5原型机，在175W功耗约束下，对于Llama3-70B模型的token生成速度达到245token/s，较传统架构提升2.8倍。这进一步验证了存内计算在生成式AI时代的战略价值。
