1. 大模型推理的瓶颈与PIM架构机遇在当今大语言模型LLM的实际部署中推理阶段特别是自回归解码过程面临着一个关键矛盾随着上下文窗口的不断扩大KV缓存Key-Value Cache的规模呈线性增长而传统计算架构中内存带宽的限制使得系统性能严重受限于数据搬运效率。以32K上下文的模型为例每个解码步骤需要访问的KV缓存可能达到GB级别这导致GPU等传统加速器的大量计算单元处于饥饿状态等待数据从显存中加载。处理内存PIM架构为解决这一困境提供了新的可能性。与传统冯·诺依曼架构不同PIM将计算单元直接嵌入内存控制器附近利用内存内部的高带宽特性实现数据就地处理。以HBM-PIM为例其每个内存通道可提供超过1TB/s的内部带宽是传统GPU显存接口带宽的5-10倍。这种架构特别适合注意力计算中频繁的矩阵-向量乘法GEMV操作因为计算与数据紧密耦合避免了数据在计算核心与内存间的反复搬运内存bank级别的并行性可以充分利用注意力计算中的天然并行性高带宽特性能够满足注意力层对内存访问的密集需求然而当前PIM架构在支持稀疏注意力时面临一个根本性矛盾现代稀疏注意力方法如Token-wise Sparsity通常以单个token为粒度进行动态选择而PIM内存的最小访问单元是完整的DRAM行通常包含多个token的数据。这种粒度不匹配导致即使只需要访问一个token系统也不得不加载并处理整行数据造成严重的计算资源浪费。2. STARC的核心设计原理2.1 语义聚类与物理映射STARC的核心创新在于通过语义聚类重构KV缓存的内存布局使动态稀疏模式与PIM的硬件特性相适配。其技术路线包含三个关键步骤在线聚类算法采用改进的K-means算法对KV对进行实时聚类距离度量使用余弦相似度。对于维度为d的key向量k_i和k_j其相似度计算为sim(k_i, k_j) (k_i · k_j) / (||k_i|| * ||k_j||)这种度量方式比传统的欧式距离更适合捕捉高维向量空间的语义关系。硬件感知的簇大小将每个簇的token数量设置为32的整数倍这与HBM3的突发传输长度8个FP16向量每个向量128维完美对齐。例如在head维度为128的LLaMA架构中一个DRAM行正好存储8个key或value向量32个token对应4个完整DRAM行。增量式聚类更新采用预填充聚类解码增量的两阶段策略预填充阶段对初始prompt的所有KV对进行全局聚类解码阶段每生成128个新token后仅对新token进行局部聚类 这种策略将聚类开销分摊到多个解码步骤实测中仅增加约3%的额外延迟。2.2 集群粒度检索机制与传统稀疏注意力逐token选择不同STARC在检索阶段操作的是整个聚类单元。其工作流程如下聚类表征为每个聚类维护一个质心向量μ初始化为该簇所有key向量的平均值快速匹配对于当前查询向量q计算其与各聚类质心的注意力分数s_j q^T μ_j / √d层级选择首先按分数降序选择完整聚类最后一个被选中的聚类可能只部分包含通过token分数排序截断并行加载将选中的聚类对应的DRAM行批量加载到PIM计算单元这种集群粒度的检索方式带来两个关键优势硬件友好每次内存访问都能充分利用PIM的行级并行性语义连贯同一聚类内的token具有相似的语义提高了每次内存访问的有效信息密度3. 硬件实现细节3.1 HBM-PIM系统适配STARC在AttAcc PIM架构上的实现需要考虑以下几个硬件约束存储布局优化// KV缓存的内存布局示例 struct { fp16 keys[CLUSTER_SIZE][HEAD_DIM]; fp16 values[CLUSTER_SIZE][HEAD_DIM]; uint32_t next_cluster_ptr; // 链表指针 } cluster_block;每个cluster_block对应一个完整的存储单元在HBM3中占据连续的物理地址空间。计算流水线Bank级并行不同bank处理不同注意力头的计算行缓冲重用同一聚类内的多个token共享部分中间计算结果提前终止当检测到某行的所有token分数均低于阈值时跳过后续计算带宽利用率def estimate_bandwidth_utilization(selected_clusters): total_tokens sum([c.size for c in selected_clusters]) accessed_rows sum([ceil(c.size / TOKENS_PER_ROW) for c in selected_clusters]) return total_tokens / (accessed_rows * TOKENS_PER_ROW)实测显示STARC的带宽利用率可达78%比传统稀疏方法提升2-3倍。3.2 能耗优化技术STARC通过三种机制显著降低能耗行激活优化减少约40%的DRAM行激活操作ACT命令计算跳过利用聚类质心的预筛选跳过约25%的低价值计算数据局部性聚类内部的高局部性减少了约35%的片上缓存冲突在HBM3的1.2V工作电压下这些优化合计带来19-27%的能耗降低。4. 实际部署考量4.1 精度-效率权衡STARC在保持模型精度方面采取了多项保障措施新鲜度保护最近生成的128个token不参与聚类直接全保留动态预算分配为每个注意力头独立调整聚类选择阈值误差补偿在softmax计算前对聚类内token分数进行温度调节在LongBench基准测试中STARC相比全KV缓存的平均精度下降仅0.8%显著优于传统页面级稀疏方法平均下降2.3%。4.2 系统集成方案在实际系统中部署STARC需要考虑以下组件混合执行引擎GPU/NPU处理QKV生成和前馈网络PIM专用于注意力计算内存管理单元维护逻辑KV缓存到物理聚类的映射表处理聚类动态增长的内存分配调度器重叠聚类更新与计算平衡不同bank的工作负载5. 性能基准测试5.1 延迟分析在LLaMA-7B模型上的测试结果显示上下文长度8K方法注意力延迟(ms/token)加速比全KV缓存12.41.0xToken-wise稀疏8.71.4xPage-wise稀疏6.22.0xSTARC5.32.3xSTARC的延迟优势随着上下文长度增加而扩大在32K上下文时达到3.1倍加速。5.2 资源利用率使用Roof-line模型分析不同方法的计算效率算术强度AI全KV缓存0.8 OPs/byteSTARC2.4 OPs/byte实际算力利用率传统稀疏35-45%STARC68-72%6. 进阶优化方向6.1 分层聚类策略为进一步提升长上下文下的效率可采用两级聚类第一级粗粒度主题聚类~1000token/簇第二级细粒度语义聚类~32token/子簇 这种分层结构可将聚类更新时间减少40%。6.2 自适应簇大小根据注意力头的特点动态调整簇大小def dynamic_cluster_size(head_dim, layer_idx): base 32 if head_dim 128: return base * 2 elif layer_idx 3: return base // 2 else: return base6.3 近存计算优化在PIM计算单元旁添加专用聚类加速器将K-means计算卸载到内存端可进一步减少5-8%的额外开销。在实际部署中我们发现STARC的效能高度依赖于KV缓存预热的充分性。建议在第一个解码步骤前进行完整的全局聚类初始化并在后续每128-256个token执行增量更新。对于需要极高吞吐的场景可以适当放宽聚类质量要求如减少K-means迭代次数这能在精度损失可控的情况下获得额外的性能提升。