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_size=1)到纯模型并行(micro_batch_size=N)的各种策略
这种设计允许系统根据当前工作负载特性动态调整并行粒度。例如,对于以解码为主的负载,可以采用较大的微批次来减少通信;而对于预填充密集的负载,则可以使用较小的微批次来提高并行度。
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,输出483token
- CNN/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推理场景中都能显著提升多芯片加速器的效率。特别是在处理动态工作负载时,其优势更加明显。框架的开源版本已经发布,欢迎社区贡献和反馈。