arXiv:2505.02533cs.DCcs.AI2025-05被引 11

动态拆分注意力头,让大模型在边缘设备低延迟运行

Large Language Model Partitioning for Low-Latency Inference at the Edge

  • 按注意力头拆分模型,实时迁移以应对资源紧张
  • 3-5设备下延迟仅比最优解高15%-20%
  • 适合边缘部署的实时大模型推理场景

基于自回归解码器的大型语言模型(LLM)逐词生成文本,每生成一个词,键值缓存(key-value caches)随之扩大,导致内存与计算需求持续上升。在资源受限的边缘环境中,传统的层级划分常引发内存溢出或高延迟。本文提出一种资源感知的Transformer分块算法,在生成过程中定期更新划分策略。该方法采用短视决策,依据设备当前资源与网络带宽即时调整分块位置。初始时将模型块分配至设备,后续根据负载动态迁移,使迁移延迟与推理延迟总和最小。通过将每个注意力头与其键值缓存共置,并允许动态迁移,实现注意力头的并行执行,显著降低推理延迟。实验表明,在3-5台设备的小规模设置中,该方法延迟仅达精确最优解的15%~20%;在大规模测试中,相比现有层级划分方法,显著提升推理速度并优化内存使用。

原文摘要 · Abstract (English)

Large Language Models (LLMs) based on autoregressive, decoder-only Transformers generate text one token at a time, where a token represents a discrete unit of text. As each newly produced token is appended to the partial output sequence, the length grows and so does the memory and compute load, due to the expanding key-value caches, which store intermediate representations of all previously generated tokens in the multi-head attention (MHA) layer. As this iterative process steadily increases memory and compute demands, layer-based partitioning in resource-constrained edge environments often results in memory overload or high inference latency. To address this and reduce inference latency, we propose a resource-aware Transformer architecture partitioning algorithm, where the partitioning decision is updated at regular intervals during token generation. The approach is myopic in that it is based on instantaneous information about device resource availability and network link bandwidths. When first executed, the algorithm places blocks on devices, and in later executions, it migrates these blocks among devices so that the sum of migration delay and inference delay remains low. Our approach partitions the decoder at the attention head level, co-locating each attention head with its key-value cache and allowing dynamic migrations whenever resources become tight. By allocating different attention heads to different devices, we exploit parallel execution of attention heads and thus achieve substantial reductions in inference delays. Our experiments show that in small-scale settings (3-5 devices), the proposed method achieves within 15 to 20 percent of an exact optimal solver's latency, while in larger-scale tests it achieves notable improvements in inference speed and memory usage compared to state-of-the-art layer-based partitioning approaches.

边缘计算大模型推理动态分区

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。