Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence
  Lengths in Large Language Models
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.
659e09c4c474a955d4e840ac
Zhen Qin
admin_assigned
2024-01-10T09:03:52.712Z
642e63a53c2cf43f6d6dc5ce
/avatars/dfd78c8d55485c22be6e616670a633e5.svg
zhenqin
user
Doreamonzzz
659e09c4c474a955d4e840ad
Weigao Sun
admin_assigned
2024-01-10T09:02:24.676Z
6246bb33da617c00b48e4d92
/avatars/0304a9f6eb7f5dee4d933d03222f94e9.svg
Weigao Sun
user
weigao266
659e09c4c474a955d4e840ae
Dong Li
admin_assigned
2024-01-10T09:04:16.624Z
61cd713f78c1c6b489a231fc
/avatars/af261b85bbd10a1182372ffc459640b8.svg
Dong Li
user
liddalidd
659e09c4c474a955d4e840af
Xuyang Shen
admin_assigned
2024-01-10T09:03:01.800Z
642e4d4d6748dd4f8eeb7732
/avatars/fd911e9143d1a7aedd21a7d611543fcc.svg
Xuyang Shen
user
Ryan1122
659e09c4c474a955d4e840b0
Weixuan Sun
admin_assigned
2024-01-10T09:03:27.460Z
640c4d2d536d9fe0f000b209
/avatars/419e59b78f7fc7f6b3ac343398d11ec1.svg
Weixuan Sun
user
weixuansun
659e09c4c474a955d4e840b1
Yiran Zhong
claimed_verified
2024-01-10T08:38:44.072Z
64c525e4d68946edad6c7067
/avatars/1b108661634af602717a4ab4b66a151f.svg
Yiran Zhong
user
IanZhong
Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence
  Lengths in Large Language Models
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.
Lightning Attention-2 efficiently processes sequences of unlimited length using a tiling technique, achieving consistent speed and outperforming conventional attention mechanisms.
linear attention softmax attention cumulative summation intra-block inter-block attention kernel tricks tiling technique Triton IO-aware hardware-friendly