面向长上下文 LLM 加速的打包-预取调度架构与超大容量芯片内存储器
硬件体系结构
2025-08-13 v1 新兴技术
摘要
长上下文大语言模型（LLM）推理随上下文长度而增长注意力计算，主要因 KV 缓存传输开销饱和了高带宽存储器（HBM）。虽然预取技术通过提前获取 KV 数据来缓解缓存未命中，但其空间和时间收益为新机会提供了利用。本工作提出一种打包-预取调度架构，配合单体 3D（M3D）后段工艺（BEOL）兼容的嵌入式存储器，拥有超大容量的芯片内存储器，以加速长上下文 LLM 推理。我们的优化显示，在使用类 TPUv6e 硬件并额外增加 512MB BEOL 存储器后，Llama3.1-8B 上实现 8.06 倍解码加速和 1.83 倍整体延迟降低。在 TPU 类架构上对多请求工作负载的评估显示，与仅使用打包的方法相比，Llama3.1-8B 和 Llama3.1-70B 模型实现 1.7 倍至 2.4 倍的吞吐量提升和 1.5 倍至 2.4 倍的 HBM 带宽降低。通过打包、预取和 BEOL 存储器的协同设计，我们的方案缓解了 HBM 约束，实现了高效的长上下文 LLM 推理。
引用
@article{arxiv.2508.08457,
title = {Architecting Long-Context LLM Acceleration with Packing-Prefetch Scheduler and Ultra-Large Capacity On-Chip Memories},
author = {Ming-Yen Lee and Faaiq Waqar and Hanchen Yang and Muhammed Ahosan Ul Karim and Harsono Simka and Shimeng Yu},
journal= {arXiv preprint arXiv:2508.08457},
year = {2025}
}
备注
7 pages, 8 figures, 2 tables