SwiftKV：基于知识保护模型转换的快速预填充优化推理
机器学习
2025-06-03 v3 人工智能
计算与语言
摘要
LLM 推理在企业应用中，如摘要生成、RAG 和代码生成，通常观察到提示比生成长度长得多，导致高预填充成本和响应延迟。我们提出了 SwiftKV，一种针对减少提示 token 预填充计算量（以 FLOPs 为单位）且保持高生成质量的模型转换和蒸馏方法。首先，SwiftKV 使用后续层的 KV 缓存填充早期层的输出，使提示 token 可跳过这些后续层。其次，SwiftKV 采用轻量级的知识保护蒸馏程序，可最小化精度损失地适配现有 LLM。最后，SwiftKV 可自然集成 KV 缓存压缩，以提高低内存场景下的推理性能。我们的全面实验表明，SwiftKV 在多个 LLM 系列中可有效将预填充计算量降低 25-50%，且质量损耗最小。在端到端推理服务中，SwiftKV 实现了最高可达 2 倍的聚合吞吐量和 60% 的输出 token 时间降低。其标准化推理吞吐量可达 560 TFlops/GPU，对应 Llama-3.1-70B 可达 16K token/s。SwiftKV 已开源于 https://github.com/snowflakedb/arctictraining。
引用
@article{arxiv.2410.03960,
title = {SwiftKV: Fast Prefill-Optimized Inference with Knowledge-Preserving Model Transformation},
author = {Aurick Qiao and Zhewei Yao and Samyam Rajbhandari and Yuxiong He},
journal= {arXiv preprint arXiv:2410.03960},
year = {2025}
}