完全开源的语言模型学习记录--TrilinearCIM架构

https://arxiv.org/pdf/2604.07628
Trilinear Compute-in-Memory Architecture for Energy-Efficient Transformer Acceleration

一、一段话总结

本文提出TrilinearCIM架构，基于双栅铁电场效应晶体管（DG‑FeFET），通过背栅调制实现三操作数乘累加，在无需运行时非易失性存储器（NVM）重编程的情况下完成Transformer自注意力全流程计算，在BERT‑base与ViT‑base上验证，相比传统FeFET CIM实现最高46.6%能耗降低、20.4%延迟降低，9项GLUE任务中7项精度更优，是首个纯NVM核内完成Transformer注意力计算的架构。

二、思维导图

## **研究背景** - Transformer自注意力：O(N²)复杂度、动态Q/K/V - 传统CIM局限：双操作数、需反复NVM写入 - 现有方案：混合CMOS/降写入但未根除 ## **核心创新** - 器件：DG‑FeFET（顶栅存权重、背栅动态调制） - 算子：三操作数MAC（A·B·C） - 数据流：融合投影与注意力、消除中间张量存储 ## **架构设计** - 层级：芯片→Tile→PE→SubArray - 阵列：无选择器结构、两种三线性配置 - 数字单元：Softmax/LayerNorm/GELU专用SFU ## **评估框架** - TransCIM：基于NeuroSim扩展、支持精度/PPA建模 - 配置：7nm CMOS+22nm FeFET、8bit输入权重 ## **实验结果** - 精度：GLUE 7/9任务优于双线性CIM - PPA：能耗-46.6%、延迟-20.4%、面积+37.3% - 消融：1b/6b最优精度、32×32最优延迟 ## **优势与局限** - 优势：零写入、高能效、长序列友好 - 局限：ViT精度下降、需硬件验证

三、详细总结

1. 研究动机与问题

• Transformer自注意力生成动态Q/K/V矩阵，传统CIM需反复NVM重编程，导致延迟高、能耗大、耐久性差。
• FeFET读写不对称：写入延迟_{50ns、读取}10ns；写入能耗亚pJ级、读取fJ级。
• BERT‑base单次推理需**~75.5M次写入**，严重限制器件寿命。

2. 核心技术方案

（1）器件基础：DG‑FeFET

• 顶栅（TG）：非易失性存储静态权重。
• 背栅（BG）：挥发性电压调制沟道电导，提供第三操作数通路。
• 电导特性：G D S ( V B G ) ≈ G 0 ⋅ ( 1 + η B G ⋅ V B G ) G_{DS}(V_{BG})≈G_0·(1+η_{BG}·V_{BG})GDS​(VBG​)≈G0​⋅(1+ηBG​⋅VBG​)，工作区间G 0 ∈ [ 29 , 69 ] μ S G_0∈[29,69]μSG0​∈[29,69]μS，η ˉ B G = 0.157 V − 1 \bar{η}_{BG}=0.157V^{-1}ηˉ​BG​=0.157V−1。

（2）三线性CIM原语

• 实现三操作数乘累加（Y=A·B·C），替代传统双操作数（输入·权重）。
• 权重一次性编程，动态操作数通过背栅电压传入，全程无NVM写入。

（3）注意力数据流

1. 缩放Query生成：R 1 = X ⋅ W Q T / d k R_1=X·W_Q^T/\sqrt{d_k}R1​=X⋅WQT​/dk​​
2. 注意力分数合成：R 2 = R 1 ⋅ W K ⋅ X T R_2=R_1·W_K·X^TR2​=R1​⋅WK​⋅XT
3. Value聚合：R e s u l t = S c o r e ⋅ X ⋅ W V T Result=Score·X·W_V^TResult=Score⋅X⋅WVT​

• 优势：消除中间Q/K/V存储，缓冲区需求降低**~3倍**。

（4）硬件架构

• 层级：芯片级（2×2 Tile）→Tile级（2×2 PE）→PE级（2×2 Array）。
• 阵列：无选择器DG‑FeFET交叉阵列，列级背栅驱动与DAC。
• 数字单元：专用SFU执行Softmax、LayerNorm、GELU。

3. 评估与结果

（1）实验配置

• 模型：BERT‑base、ViT‑base；数据集：GLUE、ImageNet/CIFAR。
• 工艺：7nm CMOS + 22nm FeFET；阵列：64×64；精度：8bit输入/权重。

（2）精度结果

（3）PPA结果（BERT‑base）

（4）消融实验

• 子阵列：32×32延迟**-40.9%**，64×64能效更高。
• 精度：1b/6b为最优精度点，能耗-37.5%、延迟-26.0%、面积+32.4%。
• 序列长度：越长，三线性零写入优势越显著。

4. 贡献与结论

• 首次实现纯NVM核内Transformer注意力计算，零运行时重编程。
• 能效、延迟显著提升，面积开销可控。
• 更适合NLP长序列场景，视觉任务需算法优化。

四、关键问题与答案

1. TrilinearCIM的核心创新点是什么？
   答：核心创新是基于DG‑FeFET的三操作数CIM原语，利用背栅调制提供第三操作数通路，将静态权重存于顶栅、动态操作数以背栅电压传入，完全消除Transformer注意力推理时的NVM写入，同时融合投影与注意力步骤，降低缓冲区需求约3倍。

2. TrilinearCIM相比传统CIM的性能收益与代价是什么？
   答：收益为最高46.6%能耗降低、20.4%延迟降低、TOPS/W提升39.2%，9项GLUE任务7项精度更优；代价是37.3%面积增加，ViT视觉任务精度因背栅DAC量化有所下降。

3. 为何TrilinearCIM在NLP任务表现更好，而视觉任务稍弱？
   答：NLP任务对噪声鲁棒性强，小扰动不改变注意力结果；ViT注意力分布稀疏、存在高幅值峰值，背栅DAC的均匀量化会扭曲关键峰值，且ViT激活通道方差更大，放大量化误差，导致精度下降。

》实战测试，随着序列长度以及隐藏层维度的升高，虽然内存下降，但是时耗明显增高

importnumpyasnpimporttime# -----------------------------------------------------------------------------# 优化版：传统注意力（高维 + 长序列 专用）# -----------------------------------------------------------------------------deftraditional_attention_fast(X,Wq,Wk,Wv,d_k,block=512):N,d=X.shape out=np.zeros((N,d_k),dtype=np.float16)foriinrange(0,N,block):Xi=X[i:i+block]Qi=Xi @ Wq.T Ki=X @ Wk.T Vi=X @ Wv.T score=(Qi @ Ki.T)/np.sqrt(d_k)attn=np.exp(score)/np.sum(np.exp(score),axis=-1,keepdims=True)out[i:i+block]=attn @ Vireturnout# -----------------------------------------------------------------------------# 极速版：TrilinearCIM 注意力（高维最强，不生成Q/K/V）# -----------------------------------------------------------------------------deftrilinear_attention_fast(X,Wq,Wk,Wv,d_k,block=512):N,d=X.shape out=np.zeros((N,d_k),dtype=np.float16)foriinrange(0,N,block):Xi=X[i:i+block]R1i=(Xi @ Wq.T)/np.sqrt(d_k)score=R1i @ Wk @ X.T attn=np.exp(score)/np.sum(np.exp(score),axis=-1,keepdims=True)out[i:i+block]=attn @ X @ Wv.Treturnout# -----------------------------------------------------------------------------# 测试：长序列 1w + 高维度 768/1024# -----------------------------------------------------------------------------if__name__=="__main__":SEQ_LEN=10000HIDDEN_DIM=128# 你可以改成 512/768/1024 测试D_K=64print(f"序列长度：{SEQ_LEN}")print(f"隐藏维度：{HIDDEN_DIM}")# 数据X=np.random.randn(SEQ_LEN,HIDDEN_DIM).astype(np.float16)Wq=np.random.randn(D_K,HIDDEN_DIM).astype(np.float16)Wk=np.random.randn(D_K,HIDDEN_DIM).astype(np.float16)Wv=np.random.randn(D_K,HIDDEN_DIM).astype(np.float16)# 传统t0=time.time()out_trad=traditional_attention_fast(X,Wq,Wk,Wv,D_K)t1=time.time()# 三线性out_tri=trilinear_attention_fast(X,Wq,Wk,Wv,D_K)t2=time.time()# 对比print(f"\n传统耗时：{t1-t0:.2f}s")print(f"三线性耗时：{t2-t1:.2f}s")print(f"🚀 三线性 加速比：{(t1-t0)/(t2-t1):.2f}x")

# cpu测试结果传统耗时：126.50s 三线性耗时：796.03s 🚀 三线性 加速比：0.16x

• gpu

importtorchimporttime# 检查GPUdevice=torch.device("cuda"iftorch.cuda.is_available()else"cpu")print("使用设备:",device)# -----------------------------------------------------------------------------# 传统注意力（GPU）# -----------------------------------------------------------------------------deftraditional_attention_gpu(X,Wq,Wk,Wv,d_k):Q=X @ Wq.T K=X @ Wk.T V=X @ Wv.T score=(Q @ K.T)/torch.sqrt(torch.tensor(d_k,dtype=X.dtype))attn=torch.softmax(score,dim=-1)out=attn @ Vreturnout# -----------------------------------------------------------------------------# TrilinearCIM 三线性注意力（GPU，无Q/K/V中间张量）# -----------------------------------------------------------------------------deftrilinear_attention_gpu(X,Wq,Wk,Wv,d_k):R1=(X @ Wq.T)/torch.sqrt(torch.tensor(d_k,dtype=X.dtype))score=R1 @ Wk @ X.T attn=torch.softmax(score,dim=-1)out=attn @ X @ Wv.Treturnout# -----------------------------------------------------------------------------# 超长篇 + 高维度 测试# -----------------------------------------------------------------------------if__name__=="__main__":# 长序列 + 高维度SEQ_LEN=10000HIDDEN_DIM=2048# BERT-base 维度，可改 1024D_K=64# 数据放到 GPUX=torch.randn(SEQ_LEN,HIDDEN_DIM,device=device,dtype=torch.float16)Wq=torch.randn(D_K,HIDDEN_DIM,device=device,dtype=torch.float16)Wk=torch.randn(D_K,HIDDEN_DIM,device=device,dtype=torch.float16)Wv=torch.randn(D_K,HIDDEN_DIM,device=device,dtype=torch.float16)torch.cuda.synchronize()t0=time.time()# 传统out_trad=traditional_attention_gpu(X,Wq,Wk,Wv,D_K)torch.cuda.synchronize()t1=time.time()# 三线性out_tri=trilinear_attention_gpu(X,Wq,Wk,Wv,D_K)torch.cuda.synchronize()t2=time.time()# 误差abs_err=torch.abs(out_trad-out_tri).mean().item()rel_err=abs_err/torch.abs(out_trad).mean().item()# 输出结果print(f"\n📊 序列长度{SEQ_LEN}| 维度{HIDDEN_DIM}| GPU 对比")print(f"传统耗时 :{t1-t0:.3f}s")print(f"三线性耗时 :{t2-t1:.3f}s")print(f"🚀 加速比 :{((t1-t0)/(t2-t1)):.2f}x")print(f"平均绝对误差:{abs_err:.6f}")print(f"平均相对误差:{rel_err:.2%}")

序列长度10000|维度2048|GPU 对比 传统耗时:0.001s 三线性耗时:0.010s 加速比:0.13x 平均绝对误差: nan 平均相对误差: nan%