frozenleaves/vllm
1
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2025-10-20 14:53:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

[BugFix] Fix nightly MLA failure (FA2 + MLA chunked prefill, i.e. V1, producing bad results) (#15492)

Browse Source 
Signed-off-by: LucasWilkinson <lwilkinson@neuralmagic.com>
This commit is contained in:
Lucas Wilkinson
2025-03-25 23:33:22 -04:00
committed by GitHub
parent 23114d3364
commit 33437bc6e7
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
vllm/attention/ops/triton_merge_attn_states.py
View File

@ -54,6 +54,15 @@ def merge_attn_states_kernel(
p_lse = tl.load(prefix_lse + head_idx * num_tokens + token_idx)
s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx)
# FA2 and FA3 have different behavior for when the sum-exp is 0, this namely
# arises with 0 len seqlens. FA3 returns -inf here while FA2 returns inf.
# If we see an inf assume FA2 and convert inf to -inf for consistency
# and correctness. Inf generally doesn't make sense in this context outside
# of undefined-behavior/FA2-case, so I think this a safe assumption.
p_lse = float('-inf') if p_lse == float('inf') else p_lse
s_lse = float('-inf') if s_lse == float('inf') else s_lse
max_lse = tl.maximum(p_lse, s_lse)
p_lse = p_lse - max_lse
s_lse = s_lse - max_lse