github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 18:40:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Kernel/Fix] Revise flash attention triton kernel API and add benchmark (#5301)

Browse Source 
* fix decoding kernel pytest

* revise and add triton context attn benchmark
This commit is contained in:
Yuanheng Zhao
2024-01-23 17:16:02 +08:00
committed by GitHub
parent 8e606ecc7e
commit 3da9993b0d
5 changed files with 116 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/inference/modeling/layers/attention.py
View File

@@ -87,7 +87,7 @@ class PagedAttention:
Transform 1D no_pad tensor into 2D padded tensor with shape [bsz,seq_len,num_heads,head_size]
"""
bsz = len(seq_lengths)
padded_tensor = torch.zeros(bsz, max_seq_len, num_heads, head_size)
padded_tensor = torch.zeros(bsz, max_seq_len, num_heads, head_size, dtype=tensor.dtype)
token_idx = 0
for i, seq_len in enumerate(seq_lengths):