Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[bloom] Add kv cache support for flash attention & fix bugs #7735

Merged
merged 12 commits into from
Dec 29, 2023
Merged

[bloom] Add kv cache support for flash attention & fix bugs #7735

Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
887a201
Add kv cache support for flash attention
w5688414 Dec 27, 2023
a899fce
Update chatglm flash attention version check
w5688414 Dec 28, 2023
3cbd143
Add test for flash attention
w5688414 Dec 28, 2023
8b7bea1
Fix unitest bug
w5688414 Dec 28, 2023
fb193b5
Add flash attention to predictor
w5688414 Dec 28, 2023
d992d42
Add flash attention2
w5688414 Dec 28, 2023
02f07e9
Add flash attention unitests
w5688414 Dec 28, 2023
0734019
Merge branch 'develop' of https://github.com/PaddlePaddle/PaddleNLP i…
w5688414 Dec 28, 2023
60b656e
fix prefix decoder
w5688414 Dec 28, 2023
731b551
remove unused comments
w5688414 Dec 29, 2023
6639388
Update unitest
w5688414 Dec 29, 2023
a5f3286
Update unitest
w5688414 Dec 29, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 15 additions & 30 deletions paddlenlp/transformers/bloom/modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -378,6 +378,20 @@ def forward(
(query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)

batch_size, q_length, _, _ = query_layer.shape

if layer_past is not None:
past_key, past_value = layer_past
# concatenate along seq_length dimension:
# - key: [batch_size, self.num_heads, head_dim, kv_length]
# - value: [batch_size, self.num_heads, kv_length, head_dim]
key_layer = paddle.concat((past_key, key_layer), axis=1)
value_layer = paddle.concat((past_value, value_layer), axis=1)

if use_cache is True:
present = (key_layer, value_layer)
else:
present = None

version = paddle.version.full_version
version_check = True
if version != "0.0.0" and version <= "2.5.2":
Expand All @@ -404,39 +418,10 @@ def forward(
attn_output = attn_output.reshape([attn_output.shape[0], attn_output.shape[1], -1])
output_tensor = self.dense(attn_output)

query_layer = query_layer.transpose([0, 2, 1, 3])
key_layer = key_layer.transpose([0, 2, 3, 1])
value_layer = value_layer.transpose([0, 2, 1, 3])
if layer_past is not None:
past_key, past_value = layer_past
# concatenate along seq_length dimension:
# - key: [batch_size, self.num_heads, head_dim, kv_length]
# - value: [batch_size, self.num_heads, kv_length, head_dim]
key_layer = paddle.concat((past_key, key_layer), axis=3)
value_layer = paddle.concat((past_value, value_layer), axis=2)

if use_cache:
present = (key_layer, value_layer)
else:
present = None
else:

query_layer = query_layer.transpose([0, 2, 1, 3])
key_layer = key_layer.transpose([0, 2, 3, 1])
value_layer = value_layer.transpose([0, 2, 1, 3])
if layer_past is not None:
past_key, past_value = layer_past
# concatenate along seq_length dimension:
# - key: [batch_size, self.num_heads, head_dim, kv_length]
# - value: [batch_size, self.num_heads, kv_length, head_dim]
key_layer = paddle.concat((past_key, key_layer), axis=3)
value_layer = paddle.concat((past_value, value_layer), axis=2)

if use_cache is True:
present = (key_layer, value_layer)
else:
present = None

_, _, _, kv_length = key_layer.shape

query_layer = query_layer.reshape([batch_size * self.num_heads, q_length, self.head_dim])
Expand Down Expand Up @@ -949,7 +934,7 @@ def forward(
seq_length_with_past = seq_length
past_key_values_length = 0
if past_key_values[0] is not None:
past_key_values_length = past_key_values[0][0].shape[3]
past_key_values_length = past_key_values[0][0].shape[1]
seq_length_with_past = seq_length_with_past + past_key_values_length

if attention_mask is None:
Expand Down