arXiv:2410.03960cs.LGcs.AI2024-10EMNLP被引 18

通过知识保持的模型转换,显著降低大模型预填充计算开销。

SwiftKV: Fast Prefill-Optimized Inference with Knowledge-Preserving Model Transformation

  • 用前层输出预填充后层的键值缓存,跳过冗余计算。
  • 预填充计算减少25%-50%,生成质量几乎无损。
  • 适合高并发企业级推理场景,如摘要与代码生成。

企业级大模型应用(如摘要、RAG、代码生成)通常具有长提示词和短生成内容,导致预填充阶段计算成本高、响应延迟大。本文提出SwiftKV,一种面向快速预填充的模型转换与蒸馏方法,在大幅降低提示词处理的浮点运算量(FLOPs)的同时保持高质量生成能力。首先,利用早期层输出预填充后期层的键值缓存(KV cache),使提示词无需经过后续层计算;其次,采用轻量级知识保持蒸馏机制,可对现有LLM进行改造,精度损失极小;第三,天然支持键值缓存压缩,在低内存环境下进一步提升性能。实验表明,SwiftKV在多个大模型家族上实现25%-50%的预填充计算降低,端到端推理中吞吐量最高提升2倍,每输出令牌耗时降低60%。其归一化推理吞吐达560 TFlops/GPU,对应Llama-3.1-70B模型可实现16,000 tokens/s。项目已开源:https://github.com/snowflakedb/arctictraining。

原文摘要 · Abstract (English)

LLM inference for enterprise applications, such as summarization, RAG, and code-generation, typically observe much longer prompt than generations, leading to high prefill cost and response latency. We present SwiftKV, a novel model transformation and distillation procedure targeted at reducing the prefill compute (in FLOPs) of prompt tokens while preserving high generation quality. First, SwiftKV prefills later layers' KV cache using an earlier layer's output, allowing prompt tokens to skip those later layers. Second, SwiftKV employs a lightweight knowledge-preserving distillation procedure that can adapt existing LLMs with minimal accuracy impact. Third, SwiftKV can naturally incorporate KV cache compression to improve inference performance in low-memory scenarios. Our comprehensive experiments show that SwiftKV can effectively reduce prefill computation by 25-50% across several LLM families while incurring minimum quality degradation. In the end-to-end inference serving, SwiftKV realizes up to 2x higher aggregate throughput and 60% lower time per output token. It can achieve a staggering 560 TFlops/GPU of normalized inference throughput, which translates to 16K tokens/s for Llama-3.1-70B. SwiftKV is open-sourced at https://github.com/snowflakedb/arctictraining.

推理优化模型蒸馏大模型

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