STARC架构：优化LLM推理的PIM与稀疏注意力融合方案

1. STARC技术背景与核心挑战在大型语言模型LLM推理过程中注意力机制的计算开销随着上下文长度呈平方级增长成为系统性能的主要瓶颈。传统解决方案主要沿着两个方向演进一是基于硬件的内存计算PIM架构通过将计算单元嵌入内存阵列来缓解带宽压力二是算法层面的稀疏注意力机制通过动态选择关键token减少计算量。然而当这两种技术结合时却产生了新的架构性矛盾。PIM架构的内存访问具有显著的行粒度特性。以HBM2存储器为例单次行激活可传输1024字节数据对应64个FP16数值但实际稀疏注意力可能只需要其中的几个有效token。我们的实测数据显示在LLAMA-7B模型的8192上下文长度下传统token-wise稀疏方法会导致93%的PIM行带宽被浪费。更棘手的是这种细粒度随机访问会引发频繁的行缓冲冲突row buffer conflict使得实际延迟比理论值高出2.8倍。现有解决方案存在明显缺陷页面级稀疏Page-wise虽然对齐PIM行粒度但强制整页选择导致准确率下降17-23%混合精度计算如PIM-LLM方案无法解决访问不规则性问题动态调度策略如PAPI引入额外元数据开销使稀疏收益被抵消关键发现在2048-8192的典型上下文范围内注意力层的能量消耗占系统总能耗的58%-72%其中超过40%来自无效数据的传输和缓冲2. STARC架构设计原理2.1 聚类驱动的KV存储布局STARC的核心创新在于将语义相似的key-value对聚类后连续存储。具体实现分为三个关键步骤在线聚类引擎采用改进的k-means算法以余弦相似度为距离度量动态调整聚类中心数k max(8, context_length/256)每处理512个token触发增量聚类耗时仅占解码时间的3.2%内存映射策略def remap_kv_cache(cluster_labels, kv_data): clustered_data defaultdict(list) for label, kv in zip(cluster_labels, kv_data): clustered_data[label].append(kv) # 按PIM行边界对齐 for label in clustered_data: pad_size ROW_SIZE - (len(clustered_data[label]) % ROW_SIZE) clustered_data[label].extend([zero_kv] * pad_size) return clustered_data该策略确保同一聚类的KV对位于连续物理地址每个内存行包含同聚类多个token保留原始token顺序的元数据索引查询感知的预取机制 当新query到来时计算其与各聚类中心的相似度按相似度降序预取聚类数据设置相似度阈值θ0.6仅加载相关度高的聚类2.2 硬件友好的执行流程与传统方案的对比以处理8192上下文为例步骤Full KVToken-wiseSTARC内存访问次数128042089有效数据利用率100%18%73%行缓冲冲突率12%67%9%元数据开销(字节)0327688192执行时序优化体现在并行加载阶段利用PIM的bank级并行性同时加载多个聚类流水线设计当第一个聚类数据到达时即开始计算与后续数据传输重叠动态精度调整对低相似度聚类使用8bit计算关键聚类保持FP163. 实现细节与性能调优3.1 内存访问优化技巧我们发现了几个关键优化点子行激活通过DRAM命令组合实现256B粒度的部分行读取Bank交错存储将同一聚类数据分散到不同bank提升并行度温度感知调度优先访问物理距离近的PIM单元降低信号延迟实测显示这些优化带来额外11%的延迟改善# 内存访问模式对比 (perf stat结果) Baseline: 3.2M LLC-load-misses, 12.4% stalled-cycles-frontend STARC: 1.7M LLC-load-misses, 6.3% stalled-cycles-frontend3.2 精度保持策略为避免聚类带来的准确率损失采用三重保障机制边界token保留每聚类强制保留最近5%的token时序局部性重要性重加权对压缩后的token应用权重补偿因子w_i w_i \cdot (1 \frac{entropy(K_i)}{max\_entropy})动态回退当检测到连续3次低置信度预测时临时切换至token-wise模式在PG-19长文本测试集上的表现方法准确率速度(tokens/s)Full KV72.3%42Token-wise70.1%68Page-wise53.8%91STARC71.6%834. 实际部署经验4.1 硬件适配方案在不同PIM架构上的实现差异平台修改点性能增益HBM-PIM定制precharge策略27%DDR5-PIM调整Burst Length为819%3D Stacked优化TSV连接调度31%重要提示在美光DDR5-PIM平台上需关闭Bank Group Interleaving否则会导致23%的性能回退4.2 典型问题排查我们总结的故障排查清单精度异常下降检查聚类中心更新频率验证相似度阈值θ是否过小监控边界token保留比例性能不达预期# 使用pmu-tools检测 ./pmu-tools/ocperf.py stat -e dram_controller/act_cmd/,dram_controller/pre_cmd/ -a -- sleep 5理想情况下ACT/PRE命令比应接近1:1若PRE过多说明行缓冲利用率低内存溢出调整聚类数k与上下文长度的关系启用动态压缩zstd -3实时压缩低重要性聚类5. 扩展应用场景STARC技术可延伸至视觉Transformer对图像patch进行空间聚类多模态模型跨模态联合聚类文本视觉token边缘设备结合PIM-NVM实现低功耗推理一个有趣的发现在代码补全任务中将聚类维度从语义相似度改为语法树深度可获得额外7%的速度提升。这提示我们针对不同任务特性调整聚类策略的重要性。该技术栈的演进方向包括与MoE架构结合实现专家选择的硬件加速适应可变上下文窗口的训练过程开发编译器自动优化策略