xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-31 08:47:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix triton compilation issue (#3984)

Browse Source 
Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
Bellk17 2024-04-12 16:41:26 -07:00 committed by GitHub
parent fbb9d9eef4
commit d04973ad54
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/attention/ops/triton_flash_attention.py
View File

@ -415,7 +415,11 @@ def attn_fwd(
return return
is_mqa = hq != hk is_mqa = hq != hk
off_h_k = off_h_q % hk if is_mqa else off_h_q if is_mqa: # noqa: SIM108
off_h_k = off_h_q % hk
else:
off_h_k = off_h_q
n_extra_tokens = 0 n_extra_tokens = 0
if seqlen_k < BLOCK_N: if seqlen_k < BLOCK_N:
n_extra_tokens = BLOCK_N - seqlen_k n_extra_tokens = BLOCK_N - seqlen_k