DUET: Disaggregated Hybrid Mamba-Transformer LLMs with Prefill and Decode-Specific Packages

TL;DR

DUET通过专用包加速Mamba-Transformer模型的预填充和解码阶段,实现4倍速度提升。

cs.AR 🔴 高级 2026-03-17 19 次浏览
Alish Kanani Sangwan Lee Han Lyu Jiahao Lin Jaehyun Park Umit Y. Ogras
大语言模型 加速器 混合模型 Mamba Transformer

核心发现

方法论

DUET框架通过将预填充和解码阶段分配给专用包,利用片上系统和高带宽内存实现高效加速。预填充包使用脉动阵列芯片,解码包使用向量单元阵列,支持混合Mamba和注意力层。

关键结果

  • 在Nemotron-H-56B、Zamba2-7B和Llama3-8B模型上,DUET实现了4倍的首个标记生成速度提升,1.4倍的吞吐量提升,以及1.5倍的标记间时间降低。
  • 与B200 GPU相比,DUET在所有工作负载中表现优异,显著提高了预填充和解码效率。
  • 消除了SSM递归和逐元素操作对矩阵乘法加速器的映射不佳问题。

研究意义

DUET框架在学术界和工业界具有重要意义,解决了大语言模型推理中的性能瓶颈,特别是在混合Mamba-Transformer架构中。它为未来的硬件设计提供了新的思路。

技术贡献

DUET通过硬件解耦实现了对预填充和解码阶段的专用加速,提出了可配置的微架构扩展,支持长序列SSM和标准矩阵乘法的高效处理。

新颖性

DUET首次将硬件解耦应用于混合Mamba-Transformer模型,解决了计算和内存带宽不匹配的问题。

局限性

  • DUET在处理极端长序列时可能仍面临内存带宽限制。
  • 对新型模型架构的适应性有待进一步验证。

未来方向

未来的研究方向包括优化DUET框架以支持更多类型的混合模型,探索更高效的硬件设计以进一步提高性能。

AI 总览摘要

大语言模型在推理过程中面临计算和内存带宽不匹配的问题,现有解决方案无法充分利用这两者的特性。DUET框架通过将预填充和解码阶段分配给专用硬件包,解决了这一瓶颈。预填充包使用脉动阵列芯片,解码包使用向量单元阵列,支持混合Mamba和注意力层。实验结果表明,DUET在多个模型和工作负载上实现了显著的性能提升,特别是在首个标记生成速度和吞吐量方面。尽管DUET在硬件解耦方面取得了突破,但在处理极端长序列时仍面临内存带宽限制。未来的研究将继续优化DUET框架,以支持更多类型的混合模型并提高其适应性。

深度分析

研究背景

近年来,大语言模型在自然语言处理领域取得了显著进展。然而,这些模型的推理过程仍然受到计算和内存带宽不匹配的限制,尤其是在混合Mamba-Transformer架构中。现有的解决方案通常依赖于同质GPU,这导致了资源的低效利用。

核心问题

大语言模型的推理过程分为计算密集的预填充阶段和内存带宽密集的解码阶段。现有的硬件架构无法同时满足这两种需求,导致性能瓶颈。

核心创新

DUET通过硬件解耦,将预填充和解码阶段分配给专用包。预填充包使用脉动阵列芯片,解码包使用向量单元阵列,支持混合Mamba和注意力层,解决了计算和内存带宽不匹配的问题。

方法详解

  • �� 预填充包使用脉动阵列芯片,支持大矩阵乘法和长序列SSM。
  • �� 解码包使用向量单元阵列,支持逐标记SSM和向量-矩阵乘法。
  • �� 两种架构均可在运行时配置,以支持混合模型。

实验设计

实验在Nemotron-H-56B、Zamba2-7B和Llama3-8B模型上进行,使用ArXiv、BWB、LongWriter和Chat等工作负载进行评估。DUET在所有工作负载中均优于B200 GPU。

结果分析

实验结果表明,DUET在首个标记生成速度上实现了4倍提升,吞吐量提高了1.4倍,标记间时间降低了1.5倍。

应用场景

DUET适用于需要高效推理的大语言模型,特别是在需要混合Mamba和Transformer层的应用中,如自然语言生成和理解。

局限与展望

尽管DUET在硬件解耦方面取得了突破,但在处理极端长序列时仍面临内存带宽限制。未来的研究将继续优化DUET框架,以支持更多类型的混合模型并提高其适应性。

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

想象一个工厂,预填充阶段就像原材料的准备,需要快速高效地处理大量数据。解码阶段则像是产品的组装,需要精细的操作和高效的资源利用。DUET就像是为这个工厂设计了两条专用的生产线,一条专注于快速处理原材料,另一条则专注于高效组装产品。通过这种方式,工厂的整体效率得到了显著提升。

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

想象你在玩一个需要快速反应的游戏,预填充阶段就像是你在收集道具,解码阶段则是你在使用这些道具来打败敌人。DUET就像是给你配备了两种专用装备,一种让你更快地收集道具,另一种让你更有效地使用道具。这样,你就能在游戏中表现得更好,打败更多的敌人!

术语表

DUET (双重加速器)

一种将预填充和解码阶段分配给专用硬件包的框架。

用于加速混合Mamba-Transformer模型的推理。

Mamba

一种结合状态空间模型和自注意力机制的混合架构。

在DUET中,与Transformer层交替使用。

SSM (状态空间模型)

一种用于高效线性时间计算的模型,具有固定大小的状态表示。

在DUET中用于长序列处理。

脉动阵列

一种用于高效矩阵乘法的硬件架构。

在DUET的预填充包中使用。

向量单元阵列

一种用于加速逐标记SSM和向量-矩阵乘法的硬件架构。

在DUET的解码包中使用。

开放问题 这项研究留下的未解疑问

  • 1 如何进一步优化DUET以支持更长的序列处理?
  • 2 在新型模型架构中,DUET的适应性如何?

应用场景

近期应用

自然语言生成

DUET可用于加速自然语言生成任务,特别是在需要混合模型的情况下。

远期愿景

智能助手

DUET可以用于开发更高效的智能助手,支持更复杂的对话和任务。

原文摘要

Large language models operate in distinct compute-bound prefill followed by memory bandwidth-bound decode phases. Hybrid Mamba-Transformer models inherit this asymmetry while adding state space model (SSM) recurrences and element-wise operations that map poorly to matmul-centric accelerators. This mismatch causes performance bottlenecks, showing that a homogeneous architecture cannot satisfy all requirements. We introduce DUET, a disaggregated accelerator that assigns prefill and decode phases to specialized packages. The Prefill package utilizes systolic array chiplets with off-package memory for efficient large matrix multiplications and long-sequence SSMs. The Decode package utilizes vector-unit arrays with high-bandwidth in-package memory to accelerate token-by-token SSM and vector-matrix multiplications. Both architectures are runtime-configurable to support hybrid models with mixed Mamba and attention layers. Evaluations on Nemotron-H-56B, Zamba2-7B, and Llama3-8B across four workloads show that DUET achieves 4x faster time to first token, 1.4x higher throughput, and 1.5x lower time between tokens over the B200 GPU.

cs.AR cs.DC