arXiv:2505.02533v1 类别：交叉学科 摘要：基于自回归解码器的大型语言模型（LLMs）逐个生成文本，其中每个token代表文本的一个离散单元。每当生成一个新的token并将其附加到部分输出序列时，序列的长度增加，同时内存和计算负载也随之增加，这是因为多头注意力（MHA）层中的扩展键值缓存存储了所有之前生成的token的中间表示。由于这一迭代过程不断增加内存和计算需求，在资源受限的边缘环境中进行分层分区往往会导致内存超载或高推断延迟。为了应对这个问题并减少推断延迟，我们提出了一种资源感知的Transformer架构分区算法，在生成token的过程中定期更新分区决策。该方法是基于设备当前可用资源和网络连接带宽的瞬时信息。在首次执行时，算法将块分配到各个设备上，并在后续执行中根据资源紧张情况在设备之间迁移这些块，以确保迁移延迟和推断延迟之和保持较低水平。我们的方法在注意力头级别对解码器进行分区，将每个注意力头与其键值缓存并置，并允许在资源紧张时动态迁移。通过将不同的注意力头分配到不同的设备，我们利用了注意力头并行执行的优势，从而实现了显著的推断延迟减少。我们的实验表明，在小型设置（3-5个设备）中，所提出的方法实现了完全最优解决者延迟的15%到20%左右，在大规模测试中，与最先进的分层分区方法相比，所提出的方法在推断速度和内存使用方面实现了显着改进。