squall/flash-attention
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Gen] Fix decode function not using top_p during iterative decoding

Browse Source
This commit is contained in:
Tri Dao 2023-08-26 15:14:22 -07:00
parent 847abe653c
commit 7a3bd55f1a
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
flash_attn/utils/generation.py
View File

@ -173,7 +173,7 @@ def decode(
teacher_outputs is None
or teacher_output_len <= inference_params.sequence_len_offset + 1
):
next_token = sample(logits, top_k=top_k, temperature=temperature)
next_token = sample(logits, top_k=top_k, top_p=top_p, temperature=temperature)
else:
next_token = teacher_outputs[:, inference_params.sequence_len_offset + 1]
sequences.append(next_token)