Source-linked AI summary
OneLA: Scaling Linear-Attention Decoding to Large Beams in Generative Recommendation
Xiangrui Yang, Cheng Peng, Yunfeng Zhao, Liang Zeng, Ao Hu, Jiawei Yang, Shengzhe Wang, Jingshan Lv, Xiao Liang, Chen Yang, Jiaqiang Liu, Yiming Qiu
TL;DR
Large-beam generative recommendation makes recurrent linear-attention state management challenging because beams share a prompt but diverge during decoding. OneLA represents that shared origin and divergence with compact records and ancestry tracking, achieving 1.54–2.46× end-to-end decode speedups over existing GDN execution paths.
Problem
Large-beam GR makes per-beam recurrent-state storage and traffic costly, and dynamic beam changes can force states to be copied, reconstructed, or recomputed.
Source: §2.3 · p. 3Method
OneLA stores the prompt-derived recurrent state once and represents each beam's divergence with compact transition records and lightweight ancestry.
Source: §1 · p. 2Results
Takeaways & Limitations
OneLA demonstrates the effectiveness of compact state representation for dynamic large-beam decoding.
Source: §5 · p. 7The exposition describes one request and one GDN head, omitting request and head indices for clarity.
Source: §3.1 · p. 4
Abstract
from arXiv · showhide
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.
1 Introduction
Generative recommendation uses large beams to produce hundreds of candidate items, making per-beam recurrent states a scaling challenge. OneLA shares prompt-derived state and tracks divergent transitions compactly, achieving 1.54–2.46× end-to-end decode speedups over existing GDN execution paths.
1 Introduction
- Large-beam GR makes independently stored recurrent states costly in storage and memory traffic, while beam reordering and duplication can force copying or recomputation.GR generates hundreds of candidates; conventional serving paths either persist each beam's full state or reconstruct it from its history.Source: §1 · pp. 1–2
- All beams share a prompt-derived recurrent state and diverge through short transition sequences, making full per-beam state materialization redundant.The divergent transitions can be represented as compact records, with replay depth bounded by the fixed-length SID sequence.Source: §1 · p. 2
- OneLA stores the shared post-prefill state once, replays compact transition records on required projections, and tracks beam ancestry with lightweight metadata.A fused GPU kernel reuses shared state across beams; OneLA reports no per-beam recurrent matrix materialization or writeback.Source: §1 · p. 2
- 1.54–2.46× end-to-end decode speedups over existing GDN execution paths.OneLA's evaluation reports this full-system result.Source: §1 · p. 2
2 Background and Motivation
This section describes SID-based generative recommendation and GDN, explains why existing serving approaches struggle with large dynamic beams, and presents the state-sharing and ancestry insights behind OneLA.
2.1 Generative Recommendation
- GR represents items as short sequences of semantic IDs and uses large-beam decoding to expand SID sequences and retain globally selected candidates.A long shared prompt and short generated suffixes combine with a large, dynamically branching beam.Source: §2.1 · p. 2
2.2 Gated Delta Network
- GDN compresses token history into a fixed-size recurrent state, updates it once per token, and reads out an output along the query.Its decay and correction factors control the state update; the query and key are used in the update and readout.Source: §2.2 · pp. 2–3
2.3 GDN Decoding for Generative Recommendation
- Existing serving paths either maintain a full recurrent state per beam or replay history independently, while dynamic beam changes can require state copying, reconstruction, or recomputation.The per-beam state capacity is Θ(w d_v d_k), and recurrent-state reads and writebacks grow with the number of live beams.Source: §2.3 · p. 3
2.4 Key Insight
- Beams share a prompt-derived state and diverge through compact transitions that can be replayed using two projections, while selection changes references rather than prior records.Append-only records plus lightweight ancestry avoid relocating state or records as the beam tree evolves.Source: §2.4 · p. 3
3 Design
OneLA shares each request’s post-prefill state across beams and represents their divergent decoding histories with compact transition records. Projection replay and ancestry tracking avoid per-beam state materialization, while a fused kernel reuses shared context across beams.
3 Design
- OneLA replaces independent per-beam recurrent states with a shared post-prefill state and short branch-specific transition histories.Its design combines replay of those transitions with lightweight ancestry for the evolving beam tree.Source: §3 · p. 3, §3.1 · p. 4
3.2 OneLA State Representation
- GDN Transition Records (GTRs) encode each beam’s update in O(d_v + d_k) storage, and projection replay computes the required query and key projections without reconstructing recurrent matrices.Each record stores a decay factor, update vector, and key vector; replay needs only two d_v-dimensional accumulators per beam and value head.Source: §3.2 · pp. 4–5
3.3 Dynamic Beam Ancestry Index
- Children reuse their selected parent’s ancestry indices, while new GTRs are appended to fixed physical slots and never relocated.Pruning discards references, reordering changes their order, and fan-out duplicates ancestry indices rather than historical GTR chains.Source: §3.3 · p. 5
3.4 Fused Shared-Context Attention Kernel
- The fused GDN kernel loads shared-context tiles once for groups of beams, replays ancestry on chip, and appends compact records instead of writing full recurrent states.The fused flow keeps intermediate projections on chip and avoids reconstructing intermediate state matrices.Source: §3.4 · p. 5
4 Evaluation
OneLA improves full-model decode-forward performance and reduces GDN latency, persistent recurrent-state capacity, state writes, and DRAM traffic. The evaluation tests the complete model and operator workloads against serving baselines, with results spanning multiple prompt lengths and configurations.
4.2 Full-Model Decode-Forward Performance
- OneLA reduces cumulative full-model decode-forward time by 1.54–2.46× across evaluated prompt and output lengths, outperforming both vLLM execution paths.Controlled FullState stays within 6.9–9.3% of native vLLM FullState.Source: Figure 6 · p. 6, §4.2 · p. 6
4.3 GR GDN Attention Generality
- Median GDN latency speedups are 40.5×/42.5× over vLLM/SGLang FullState, 17.5× over both ReplaySSM paths, and 41.5×/46.3× over FlashInfer/TensorRT-LLM.Across 405 OneLA cases, maximum absolute output error against FP32 FullState is 4.88 × 10^-4.Source: Figure 7 · p. 6, §4.3 · p. 6
- OneLA reduces median persistent recurrent-state capacity by up to 42.5× relative to baselines across paired workloads.The evaluated configurations span request counts, beam widths, decode lengths, and nine recurrent-state geometries.Source: Figure 7 · p. 6, §4.3 · p. 6
4.1 Methodology
- The evaluation pairs OneLA with native and controlled FullState/ReplaySSM baselines in a 0.8B Qwen3.5 full-model setup and across 405 operator shapes.The full-model setup uses four requests, beam width 256, prompt lengths of 1K/5K, and 3–7 output tokens.Source: §4.1 · p. 6, §4.2 · p. 6
4.4 Sources of Performance Improvement
- At O=3 and O=7, OneLA recurrent-state-path latency is 0.0755 and 0.308 ms, 39.7× and 35.4× lower than FullState.The corresponding vLLM ReplaySSM latencies are 1.20 and 4.06 ms; FullState latencies are 3.00 and 10.9 ms.Source: Figure 8 · p. 6, §4.4 · p. 7
- At O=3/7, OneLA uses 36.3/101 MiB of persistent state, 56.5×/20.3× below FullState, and writes 32.3/96.8 MiB, 127× below at both endpoints.The representation combines one request-level Sctx with compact branch-specific GTRs instead of per-beam full-state matrices.Source: Figure 8 · p. 6, §4.4 · p. 7
- At O=3, OneLA moves 161 MiB of DRAM data, 76.9× and 39.9× less than FullState and ReplaySSM, respectively.Its arithmetic intensity is 40.0 FLOP/byte, 45.6× and 20.6× higher than FullState and ReplaySSM.Source: Figure 8 · p. 6, §4.4 · p. 7
5 Conclusion
OneLA applies a compact shared-state representation to large-beam generative recommendation and achieves end-to-end speedups.
5 Conclusion
- OneLA achieves 1.54–2.46× end-to-end speedups and substantial operator-level gains for large-beam generative recommendation.It replaces per-beam recurrent states with a shared post-prefill state, compact transition records, and lightweight ancestry.Source: §5 · p. 7