A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 17:15 UTC
[BugFix] Fix nightly MLA failure (FA2 + MLA chunked prefill, i.e. V1, producing bad results) (#15492)
Signed-off-by: LucasWilkinson <lwilkinson@neuralmagic.com>
Lucas Wilkinson committed
Mar 25, 2025 at 23:33 UTC
33437bc6e7af316fa9ce6b6e559501ca45d9cd45
1 file changed
+9
vllm/attention/ops/triton_merge_attn_states.py
+9
@@ -54,6 +54,15 @@ def merge_attn_states_kernel(
54
55
p_lse = tl.load(prefix_lse + head_idx * num_tokens + token_idx)
56
s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx)
57
+
58
+ # FA2 and FA3 have different behavior for when the sum-exp is 0, this namely
59
+ # arises with 0 len seqlens. FA3 returns -inf here while FA2 returns inf.
60
+ # If we see an inf assume FA2 and convert inf to -inf for consistency
61
+ # and correctness. Inf generally doesn't make sense in this context outside
62
+ # of undefined-behavior/FA2-case, so I think this a safe assumption.
63
+ p_lse = float('-inf') if p_lse == float('inf') else p_lse
64
+ s_lse = float('-inf') if s_lse == float('inf') else s_lse
65
+
66
max_lse = tl.maximum(p_lse, s_lse)
67
p_lse = p_lse - max_lse
68
s_lse = s_lse - max_lse