Heterogeneous LLM Serving with General-Purpose Processing-Near-Memory for Retrieval-Based Sparse Attention

TL;DR

提出KARAT异构系统,利用通用处理近存存储提升长上下文稀疏注意效率,性能提升达6.13倍。

cs.AR 🔴 高级 2026-08-04 60 次浏览
Hyungkyu Ham Junhyeong Bae Seungheon Lee Myeongjae Jeon Gwangsun Kim
大规模语言模型 稀疏注意 异构硬件 近存处理 系统优化

核心发现

方法论

本文分析了长上下文下检索式稀疏注意的硬件瓶颈,提出一种异构解码系统,将KV缓存迁移出GPU内存。系统由GPU节点负责模型参数和密集计算,PNM(处理近存)节点负责存储KV缓存和索引,执行读取操作。设计基于四个关键需求:容量、带宽、计算能力和通用性,提出KARAT(KV缓存驻留加速器)方案,结合大容量LPDDR和通用GPNM核,支持多样算法。引入微批次调度(OFMS)和上下文感知重平衡(CMR)以优化流水线效率。实验证明在三种先进模型上,性能提升达2.09-6.13倍,能耗效率提升1.36-3.21倍。

关键结果

  • 在三种最先进模型(如GLM-5.2、MiniMax-M3、DeepSeek-V3.2)上,系统在满足服务级别目标(每个Token 100ms)下,吞吐量比GPU单一系统提升2.09到6.13倍,显著缓解容量和带宽瓶颈。
  • 通过引入KARAT硬件,支持多样稀疏注意算法,避免固定功能单元的局限性,提升算法演进的适应性。
  • 微批次调度和上下文感知重平衡显著减少流水线空泡,提升整体吞吐率和资源利用率。

研究意义

该研究突破了长上下文大模型的硬件瓶颈限制,为未来大规模LLM的高效部署提供了新思路。通过异构硬件协作,有效平衡容量、带宽和计算需求,推动大模型在实际应用中的落地,尤其适用于多轮对话和长文本理解场景。该方案不仅提升了吞吐率,还降低了能耗,为大模型的商业化部署奠定基础。其通用性和扩展性为未来多样化稀疏注意算法提供了硬件支持,推动AI硬件设计向更高性能和能效方向发展。

技术贡献

本文提出一种支持多样稀疏注意算法的通用PNM架构KARAT,结合大容量LPDDR和灵活GPNM核,突破了传统PIM/PNM固定功能单元的限制。设计实现了容量、带宽、计算和通用性四个关键需求,采用微批次调度(OFMS)和上下文感知重平衡(CMR)优化流水线效率。系统架构实现异构设备协作,显著提升长上下文下的吞吐率和能效,推动硬件支持大规模LLM的多样化算法。

新颖性

首次提出支持多样稀疏注意算法的通用PNM硬件架构,突破了固定功能单元的局限,结合大容量LPDDR存储与通用核,实现容量、带宽和计算的协同优化。引入微批次调度和上下文感知重平衡,显著提升流水线效率,满足长上下文场景的硬件需求,推动异构硬件在LLM中的应用创新。

局限性

  • 系统设计依赖高速网络连接,网络延迟可能影响微批次调度效率,尤其在极端长上下文或高并发场景下。
  • 硬件实现复杂,需集成多种存储和处理单元,制造成本和能耗较高,限制大规模推广。
  • 算法多样性带来的硬件适配挑战,未来需持续优化硬件灵活性以支持新兴稀疏注意机制。

未来方向

未来将探索更高效的微批次调度策略,减少流水线空泡;优化硬件架构以降低成本和能耗;支持更丰富的稀疏注意算法,提升系统适应性和扩展性。同时,结合AI芯片的定制设计,推动异构系统在实际大规模应用中的部署。

AI 总览摘要

随着大规模语言模型(LLM)在多领域的广泛应用,处理长上下文成为核心挑战之一。传统GPU架构在存储容量和带宽方面难以满足百万级Token的需求,导致性能瓶颈。为此,本文提出了一种异构解码系统,将KV缓存从GPU迁移到处理近存(PNM)设备,显著缓解存储和带宽压力。系统由GPU节点负责模型参数和密集计算,KARAT(KV缓存驻留加速器)节点则存储KV缓存和索引,执行读取操作。设计基于四个关键需求:容量、带宽、计算能力和通用性,提出支持多样稀疏注意算法的硬件架构。引入微批次调度(OFMS)和上下文感知重平衡(CMR)技术,有效减少流水线空泡,提升吞吐率。在三种最先进模型上的实验显示,性能提升达2.09-6.13倍,能效提升1.36-3.21倍。这一创新架构为大模型的高效部署提供了新路径,推动AI硬件向更高性能和能效迈进。未来工作将继续优化调度策略和硬件设计,支持更丰富的算法变体,满足不断增长的应用需求。

深度分析

研究背景

近年来,LLM在自然语言处理中的表现不断突破,代表性工作如GPT系列、BERT、GLaM等推动模型规模和能力的提升。然而,随着模型规模的扩大,长上下文处理成为瓶颈,传统GPU架构在存储容量和带宽方面难以满足需求。稀疏注意机制如检索式稀疏注意(Retrieval-based Sparse Attention)被引入以减轻计算和存储压力,具体算法包括GQA、DeepSeek等,显著降低KV读写量。尽管如此,存储容量、带宽和计算能力的瓶颈依然存在,限制了模型的实际部署。此前的方案多采用将KV缓存迁移到主机内存或利用PIM/PNM硬件,但多为固定功能单元,难以支持算法演进,且在高负载下效率不足。本文旨在突破这些限制,提出一种支持多样稀疏注意算法的异构硬件架构,优化长上下文下的性能表现。

核心问题

核心问题在于长上下文场景下,存储容量、带宽和计算能力的瓶颈严重制约模型性能。GPU内存有限,难以存放百万级Token的KV缓存,带宽限制导致读取延迟高。此外,固定功能的PIM/PNM硬件无法适应算法多样性,限制了算法创新和演进。如何设计一种既能提供大容量存储,又能支持多样算法的通用硬件架构,成为亟待解决的难题。系统必须在保证高吞吐率的同时,兼顾能效和扩展性,满足实际应用中的多轮对话和长文本理解需求。

核心创新

本研究的创新点主要包括:1)提出KARAT架构,将KV缓存存放在大容量LPDDR存储中,结合通用GPNM核实现多样算法支持,突破固定功能单元限制;2)设计微批次调度(OFMS)和上下文感知重平衡(CMR)技术,有效减少流水线空泡,提升吞吐率;3)实现异构设备协作,将模型参数和密集计算留在GPU,稀疏读取操作由KARAT节点完成,优化整体性能。这些创新使系统在容量、带宽、计算和通用性方面实现了平衡,满足长上下文场景的硬件需求。

方法详解

  • �� 系统架构:将解码流程划分为GPU和KARAT两个异构部分,GPU负责模型参数和密集计算,KARAT负责存储KV缓存和索引,执行读取操作。
  • �� 硬件设计:KARAT设备集成大容量LPDDR存储和通用GPNM核,支持多样稀疏算法。
  • �� 调度策略:引入OFMS,通过微批次交错隐藏通信延迟;采用CMR,平衡不同请求的Token数,减少流水线空泡。
  • �� 数据流:预填充阶段将KV缓存迁移到KARAT,解码时GPU构建索引,KARAT执行索引扫描、Gather和Attention,最后将结果返回GPU进行后续处理。
  • �� 系统优化:通过调度和数据布局,最大化设备利用率,减少通信开销,确保在长上下文下保持高吞吐。

实验设计

采用GLM-5.2、MiniMax-M3、DeepSeek-V3.2三种代表性模型,使用真实多轮对话和长文本数据集进行测试。对比GPU单一系统,评估吞吐量、能效和服务级别目标(每Token 100ms)。设置不同微批次大小,测试调度策略的效果。通过消融实验验证OFMS和CMR的贡献,分析硬件利用率和流水线效率。实验结果显示,系统在满足SLO的情况下,吞吐率提升达2.09-6.13倍,能效提升1.36-3.21倍,验证了设计的有效性。

结果分析

实验表明,所提系统在三种模型上均实现显著性能提升,最大达6.13倍,能耗效率提升达3.21倍。微调调度策略显著减少流水线空泡,提升设备利用率。支持多样稀疏算法,展现出良好的算法适应性。系统在长上下文场景下,保持低延迟和高吞吐,满足实际应用需求。与传统GPU系统相比,显著降低了存储和带宽压力,为大规模LLM部署提供了可行方案。

应用场景

该系统适用于需要处理超长文本和多轮对话的AI应用,如智能助手、内容生成和长文本理解。通过异构硬件协作,提升模型吞吐和能效,降低部署成本。未来可扩展到多模态模型和更复杂的稀疏注意机制,推动AI在实际场景中的广泛应用。

局限与展望

系统对高速网络依赖较大,网络延迟可能影响调度效率。硬件集成复杂,制造成本较高,限制大规模推广。算法多样性带来硬件适配挑战,未来需持续优化硬件灵活性和调度策略,以支持更多新兴稀疏注意算法。

通俗解读 非专业人士也能看懂

想象你在一家大型工厂里,生产线需要处理各种不同的任务。有些任务需要用到大量的原料(存储容量),而有些任务则需要快速的机械手(计算能力)来完成。传统的工厂只用一种机器(GPU),但当任务变得复杂、原料变多时,这种单一机器就变得不够用了。于是,工厂引入了一种新方案:在不同的区域放置不同的机器,比如一部分专门存放原料(KV缓存),另一部分专门做机械操作(模型计算)。这样一来,原料存储和机械操作可以同时进行,不会互相阻碍。通过智能调度(微批次调度和重平衡),工厂可以更高效地完成任务,节省时间和能源。这就像本文提出的KARAT系统,把存储和计算分开,合作完成复杂的任务,效率大大提升。

简单解释 像给14岁少年讲一样

想象你在玩一个超级复杂的拼图游戏,你需要把很多碎片拼在一起,拼完后还能快速找到你需要的那一块。以前,所有的碎片都放在一个箱子里,你每次找东西都得翻半天,超级慢。而现在,你的朋友帮你把碎片分类存放在不同的箱子里:一个箱子专门存放大块碎片(模型参数),另一个箱子存放小碎片(KV缓存和索引)。当你拼图时,你只需要从分类好的箱子里快速找到需要的碎片,然后拼在一起。这样一来,你就可以更快完成拼图,甚至可以拼出更复杂的图案。这个系统就像论文里的KARAT,把存储和拼图(计算)分开,让整个过程更快、更省力。未来,还可以让更多朋友加入,帮你拼得更快!

原文摘要

This paper presents a heterogeneous decode-phase serving system that relocates the KV cache out of GPU memory, motivated by the retrieval-based sparse attention that recent frontier LLMs adopt to serve million-token contexts. It partitions a decode step by operation type: GPU nodes hold the model weights and execute the projections and MoE layers, while processing-near-memory (PNM) nodes hold the KV cache and index keys and execute every operation that reads them. We first show that the assumptions behind prior PIM and PNM designs no longer hold for these operations, and derive four design requirements for such a node. From these requirements, we propose KARAT (KV-cache-resident Accelerator for Retrieval-based ATtention), a general-purpose PNM design that is the design point meeting all four. A KARAT device combines large LPDDR capacity with general-purpose compute sized for the retrieval indexer, serving an operational intensity beyond what PIM/PNM designs built for low-intensity GEMV target while accommodating diverse sparse attention algorithms that fixed-function units cannot support as they evolve. To reduce pipeline bubbles as the two device types alternate between micro-batches, we further propose opportunistic, fine-grained micro-batch scheduling (OFMS), which hides expert all-to-all behind the other micro-batch's GEMMs, and context-length-aware micro-batch rebalancing (CMR), which equalizes their token counts despite the variance in context length. Across three state-of-the-art models and real agentic traces, our proposed system improves throughput per TDP under a service-level objective by 2.09-6.13x over a GPU-only baseline and runs training-free sparse attention methods with 1.36-3.21x improvements.

cs.AR