Skip to content

GitLab

Unverified Commit d04973ad authored by Bellk17's avatar Bellk17 Committed by GitHub
Browse files

Fix triton compilation issue (#3984) 


Co-authored-by: default avatarWoosuk Kwon <woosuk.kwon@berkeley.edu>
parent fbb9d9ee
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 1 deletion
vllm/attention/ops/triton_flash_attention.py
View file @ d04973ad
...@@ -415,7 +415,11 @@ def attn_fwd( ...@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment