title: Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence
    Lengths in Large Language Models
summary: Linear attention is an efficient attention mechanism that has recently
  emerged as a promising alternative to conventional softmax attention. With its
  ability to process tokens in linear computational complexities, linear
  attention, in theory, can handle sequences of unlimited length without
  sacrificing speed, i.e., maintaining a constant training speed for various
  sequence lengths with a fixed memory consumption. However, due to the issue
  with cumulative summation (cumsum), current linear attention algorithms cannot
  demonstrate their theoretical advantage in a causal setting. In this paper, we
  present Lightning Attention-2, the first linear attention implementation that
  enables linear attention to realize its theoretical computational benefits. To
  achieve this, we leverage the thought of tiling, separately handling the
  intra-block and inter-block components in linear attention calculation.
  Specifically, we utilize the conventional attention computation mechanism for
  the intra-blocks and apply linear attention kernel tricks for the inter-blocks.
  A tiling technique is adopted through both forward and backward procedures to
  take full advantage of the GPU hardware. We implement our algorithm in Triton
  to make it IO-aware and hardware-friendly. Various experiments are conducted on
  different model sizes and sequence lengths. Lightning Attention-2 retains
  consistent training and inference speed regardless of input sequence length and
  is significantly faster than other attention mechanisms. The source code is
  available at https://github.com/OpenNLPLab/lightning-attention.
paper_authors._id: 659e09c4c474a955d4e840ac
paper_authors.name: Zhen Qin
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-01-10T09:03:52.712Z
paper_authors._id: 642e63a53c2cf43f6d6dc5ce
paper_authors.avatarUrl: /avatars/dfd78c8d55485c22be6e616670a633e5.svg
paper_authors.fullname: zhenqin
paper_authors.type: user
paper_authors.user: Doreamonzzz
paper_authors._id: 659e09c4c474a955d4e840ad
paper_authors.name: Weigao Sun
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-01-10T09:02:24.676Z
paper_authors._id: 6246bb33da617c00b48e4d92
paper_authors.avatarUrl: /avatars/0304a9f6eb7f5dee4d933d03222f94e9.svg
paper_authors.fullname: Weigao Sun
paper_authors.type: user
paper_authors.user: weigao266
paper_authors._id: 659e09c4c474a955d4e840ae
paper_authors.name: Dong Li
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-01-10T09:04:16.624Z
paper_authors._id: 61cd713f78c1c6b489a231fc
paper_authors.avatarUrl: /avatars/af261b85bbd10a1182372ffc459640b8.svg
paper_authors.fullname: Dong Li
paper_authors.type: user
paper_authors.user: liddalidd
paper_authors._id: 659e09c4c474a955d4e840af
paper_authors.name: Xuyang Shen
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-01-10T09:03:01.800Z
paper_authors._id: 642e4d4d6748dd4f8eeb7732
paper_authors.avatarUrl: /avatars/fd911e9143d1a7aedd21a7d611543fcc.svg
paper_authors.fullname: Xuyang Shen
paper_authors.type: user
paper_authors.user: Ryan1122
paper_authors._id: 659e09c4c474a955d4e840b0
paper_authors.name: Weixuan Sun
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-01-10T09:03:27.460Z
paper_authors._id: 640c4d2d536d9fe0f000b209
paper_authors.avatarUrl: /avatars/419e59b78f7fc7f6b3ac343398d11ec1.svg
paper_authors.fullname: Weixuan Sun
paper_authors.type: user
paper_authors.user: weixuansun
paper_authors._id: 659e09c4c474a955d4e840b1
paper_authors.name: Yiran Zhong
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-01-10T08:38:44.072Z
paper_authors._id: 64c525e4d68946edad6c7067
paper_authors.avatarUrl: /avatars/1b108661634af602717a4ab4b66a151f.svg
paper_authors.fullname: Yiran Zhong
paper_authors.type: user
paper_authors.user: IanZhong
paper_title: Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence
    Lengths in Large Language Models
paper_summary: Linear attention is an efficient attention mechanism that has recently
  emerged as a promising alternative to conventional softmax attention. With its
  ability to process tokens in linear computational complexities, linear
  attention, in theory, can handle sequences of unlimited length without
  sacrificing speed, i.e., maintaining a constant training speed for various
  sequence lengths with a fixed memory consumption. However, due to the issue
  with cumulative summation (cumsum), current linear attention algorithms cannot
  demonstrate their theoretical advantage in a causal setting. In this paper, we
  present Lightning Attention-2, the first linear attention implementation that
  enables linear attention to realize its theoretical computational benefits. To
  achieve this, we leverage the thought of tiling, separately handling the
  intra-block and inter-block components in linear attention calculation.
  Specifically, we utilize the conventional attention computation mechanism for
  the intra-blocks and apply linear attention kernel tricks for the inter-blocks.
  A tiling technique is adopted through both forward and backward procedures to
  take full advantage of the GPU hardware. We implement our algorithm in Triton
  to make it IO-aware and hardware-friendly. Various experiments are conducted on
  different model sizes and sequence lengths. Lightning Attention-2 retains
  consistent training and inference speed regardless of input sequence length and
  is significantly faster than other attention mechanisms. The source code is
  available at https://github.com/OpenNLPLab/lightning-attention.
paper_ai_summary: Lightning Attention-2 efficiently processes sequences of unlimited length using a tiling technique, achieving consistent speed and outperforming conventional attention mechanisms.
paper_ai_keywords: linear attention softmax attention cumulative summation intra-block inter-block attention kernel tricks tiling technique Triton IO-aware hardware-friendly