DUET: Disaggregated Hybrid Mamba-Transformer LLMs with Prefill and Decode-Specific Packages
This work addresses efficiency issues for AI practitioners deploying hybrid models, offering a hardware solution to improve inference performance, though it is incremental as it builds on existing accelerator designs.
The paper tackles the performance bottlenecks in hybrid Mamba-Transformer large language models caused by mismatched compute and memory requirements between prefill and decode phases, introducing DUET, a disaggregated accelerator with specialized packages that achieves 4x faster time to first token, 1.4x higher throughput, and 1.5x lower time between tokens compared to the B200 GPU.
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.