多芯片加速器动态LLM推理优化与Compass框架实践

1. 多芯片加速器与动态LLM推理的挑战在当今AI领域大语言模型(LLM)已经成为自然语言处理任务的核心驱动力。然而这些模型的庞大规模带来了前所未有的计算挑战。单个芯片的处理能力已经难以满足LLM推理的实时性要求这使得多芯片加速器架构成为必然选择。多芯片加速器通过将计算任务分配到多个处理单元上并行执行理论上可以线性提升系统吞吐量。但在实际应用中特别是在动态LLM推理场景下这种架构面临着几个关键挑战首先现实中的LLM推理服务需要同时处理不同类型的请求。典型的请求类型包括预填充(prefill)请求处理用户输入的初始阶段需要完整计算整个输入序列的注意力解码(decode)请求生成每个token的后续阶段只需计算最后一个token的注意力其次序列长度的极端变化是另一个主要挑战。在ShareGPT等实际场景中序列长度可能从几个token到数万个token不等。这种变化不仅影响单个批次内的计算模式还会导致批次间工作负载的巨大差异。提示在实际部署中预填充阶段通常占整个推理时间的20-30%但却消耗了80%以上的计算资源。如何平衡这两种请求类型的处理是优化的关键。2. 传统映射方案的局限性现有的多芯片加速器映射方案主要针对传统的CNN/Transformer模型设计它们在处理动态LLM工作负载时表现出明显的不足。这些局限性主要体现在两个方面2.1 静态批处理假设大多数现有方案基于静态批处理假设即同一批次内的所有请求类型相同所有请求的序列长度固定或变化很小计算和内存访问模式在整个批次中保持一致这种假设简化了映射问题但与现实中的LLM推理场景严重不符。现代推理服务系统如Orca和vLLM已经采用了动态批处理策略允许混合不同类型的请求和变长序列。2.2 不完整的映射空间现有方法可以大致分为两类单模型映射将整个LLM视为单一计算图无法处理批次内的多样性多模型映射将每个请求视为独立模型忽略了LLM特有的合并-拆分-再合并执行模式这两种方法都导致映射空间不完整无法充分利用多芯片架构的并行潜力。特别是在处理混合请求类型时现有方案往往会产生大量冗余计算或通信开销。3. 计算执行图映射编码方案针对上述挑战我们提出了一种创新的计算执行图映射编码方案。该方案的核心思想是将LLM工作负载建模为一个二维计算图两个维度分别是3.1 微批次维度解耦通过引入micro_batch_size参数我们实现了微批次维度的灵活划分micro_batch_size必须是批次大小N的约数每个微批次可以独立映射到不同芯片支持从纯数据并行(micro_batch_size1)到纯模型并行(micro_batch_sizeN)的各种策略这种设计允许系统根据当前工作负载特性动态调整并行粒度。例如对于以解码为主的负载可以采用较大的微批次来减少通信而对于预填充密集的负载则可以使用较小的微批次来提高并行度。3.2 层维度分割segmentation参数是一个长度为M-1的二进制向量用于控制层维度的分割segmentation[i]1表示在第i层后插入分割点全0向量表示层优先调度全1向量表示微批次优先调度混合模式可以实现更复杂的流水线并行这种灵活的分割机制使得系统能够根据模型结构和硬件特性优化数据流。例如在注意力层前后插入分割点可以减少中间结果的存储压力。3.3 子图到芯片的映射layer_to_chip矩阵将每个子图明确映射到特定芯片矩阵尺寸为(N/micro_batch_size)×M每个元素表示对应子图的目标芯片ID支持任意复杂的跨芯片通信模式通过精心设计layer_to_chip矩阵可以实现各种混合并行策略。例如可以将前几层映射到一组芯片做数据并行后几层映射到另一组芯片做模型并行。4. Compass框架设计与实现基于上述编码方案我们开发了Compass框架它由两个核心组件构成4.1 评估引擎评估引擎负责精确预测给定映射方案的性能指标包括4.1.1 延迟模型我们采用细粒度的依赖分析来计算总执行时间T_proc,l max(T_comp,l, T_DRAM,l, T_NoP,l) T_start,l max(max(T_end,l for l in Pre(l)), max(T_end,l for l in SameCore(l))) T_end,l T_start,l T_proc,l T_model max(T_end,l for all l)其中考虑了计算、DRAM访问和芯片间通信的流水线重叠。4.1.2 能耗模型总能耗是各层能耗的累加E_proc,l E_comp,l E_DRAM,l E_NoP,l E_model sum(E_proc,l for all l)我们特别关注数据访问能耗通过算法1确定何时可以避免不必要的DRAM访问。4.2 遗传算法优化引擎为了高效搜索巨大的映射空间我们设计了专门的遗传算法4.2.1 染色体表示每个个体由三部分组成micro_batch_size整数型基因segmentation二进制向量基因layer_to_chip整数矩阵基因这种表示完全对应我们的映射编码方案确保所有可能的映射都能被表达。4.2.2 遗传操作我们设计了多种变异算子来平衡探索和开发局部微调单个基因位或矩阵元素的随机变化子图级变异整行或整列的重新映射全局重组大规模的结构调整选择压力通过锦标赛选择机制动态调节避免过早收敛。5. 实际应用与性能评估我们在三种典型硬件配置上评估了Compass框架5.1 硬件配置WS架构6×6权重固定型芯片OS架构6×6输出固定型芯片HE架构3×6 WS 3×6 OS混合型所有芯片采用TSMC 12nm工艺主频1GHz配备2MB全局缓存和1024个MAC单元。5.2 工作负载场景我们测试了两种典型序列分布ShareGPT分布平均输入78token输出483tokenCNN/DM分布平均输入866token输出63token每种分布结合三种服务策略vLLM预填充优先Orca迭代级调度Chunked Prefill分块预填充5.3 性能结果对比SCAR和MOHaM等先进方案Compass实现了平均EDP降低63.12%最高达89.61%的异构架构优势对不同模型架构的良好适应性特别值得注意的是在HE混合架构上Compass能够自动发现传统方法难以找到的优化映射充分发挥了异构计算的优势。6. 实施建议与优化技巧基于我们的实践经验为实际部署提供以下建议6.1 微批次大小选择微批次大小的选择应考虑工作负载特性解码为主选较大值预填充为主选较小值芯片数量通常设为芯片数的整数倍或约数内存限制确保单个微批次能放入芯片缓存一个好的启发式是从芯片数量的1/2倍开始尝试逐步调整。6.2 分割策略优化层分割点的设置应关注计算密集型层如注意力前后的分割内存密集型层如LayerNorm单独分割保持流水线各阶段负载均衡实践中可以先在模型的关键位置设置少量分割点再逐步细化。6.3 混合并行策略有效的混合策略通常包括底部几层数据并行处理输入多样性中间层张量并行平衡计算和通信顶部几层模型并行减少参数同步Compass的遗传算法能够自动发现这种混合策略但人工先验可以加速搜索过程。7. 典型问题排查在实际部署中我们总结了以下常见问题及解决方案7.1 性能不达预期可能原因微批次大小与硬件不匹配分割点设置不合理导致流水线气泡芯片间通信成为瓶颈解决方法尝试不同的micro_batch_size值使用Compass的可视化工具分析关键路径考虑增加NoP带宽或优化通信模式7.2 内存溢出可能原因单个微批次太大中间结果保存过多权重重复存储解决方法减小micro_batch_size调整segmentation减少同时活跃的结果启用权重共享选项7.3 收敛速度慢可能原因遗传算法参数不合适搜索空间过大评估开销太高解决方法增加种群规模和迭代次数添加人工先验约束搜索空间使用简化模型进行初步搜索经过多次实际部署验证Compass框架在各类LLM推理场景中都能显著提升多芯片加速器的效率。特别是在处理动态工作负载时其优势更加明显。框架的开源版本已经发布欢迎社区贡献和反馈。