OneLA speeds up large beam decoding for generative recommendations
OneLA: Scaling Linear-Attention Decoding to Large Beams in Generative Recommendation
Artificial IntelligenceDistributed, Parallel, and Cluster ComputingInformation Retrieval
Summary
Generating recommendations often requires looking at many possible options at once, which can use a lot of computer memory and slow down the process. The authors introduce OneLA, a way to share information across all these options so the computer doesn't have to repeat work or store as much data. Their method uses a clever system to keep track of only the differences between options, making decoding faster and less memory-intensive. This helps recommendation systems produce lots of suggestions more quickly and efficiently.
What this means in practice
- •For recommendation system engineers: Increase decoding speed and reduce memory use when producing large numbers of recommendations in generative models.
- •For natural language processing teams: Optimize beam search decoding in large-scale sequence generation tasks beyond recommendations by efficiently managing attention states.
Authors
Xiangrui Yang, Cheng Peng, Yunfeng Zhao, Liang Zeng, Ao Hu, Jiawei Yang, Shengzhe Wang, Jingshan Lv, Xiao Liang, Chen Yang, Jiaqiang Liu, Yiming Qiu
Abstract
Generative recommendation (GR) relies on large-beam decoding to generate hundreds of candidate items, creating a new scaling challenge for recurrent linear attention. Existing linear attention serving systems either materialize a full recurrent state for every beam or repeatedly replay shared history, incurring substantial memory and traffic overhead. To address this, we present OneLA, a linear-attention decoding framework that exploits the shared prompt and short divergent suffixes of GR workloads. Specifically, OneLA represents all beam states using a single shared prompt-derived state and compact, append-only records of their divergent transitions. Using this representation, OneLA computes only the state information required at each decoding step, without reconstructing a full recurrent state for every beam. Furthermore, OneLA uses a lightweight ancestry index to track the transition records that make up each beam's history, allowing beams to be updated without moving or copying existing records. A fused GPU kernel further reuses the shared state across beams. Our analysis shows that OneLA achieves 1.54-2.46x end-to-end decode speedups while substantially reducing recurrent-state memory use and data movement.