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, 14:45 UTC
[Kernel] Remove redundant Exp calculations (#16123)
Signed-off-by: DefTruth <qiustudent_r@163.com>
DefTruth committed
Apr 15, 2025 at 20:58 UTC
280d62b8a2cfd456e42aa3e4f9abb59fb87124ca
1 file changed
+6
-3
vllm/attention/ops/triton_merge_attn_states.py
+6
-3
@@ -66,7 +66,10 @@ def merge_attn_states_kernel(
66
max_lse = tl.maximum(p_lse, s_lse)
67
p_lse = p_lse - max_lse
68
s_lse = s_lse - max_lse
69
- out_se = (tl.exp(p_lse) + tl.exp(s_lse))
69
+ # Will reuse precomputed Exp values for scale factor computation.
70
+ p_se = tl.exp(p_lse)
71
+ s_se = tl.exp(s_lse)
72
+ out_se = (p_se + s_se)
73
74
if OUTPUT_LSE:
75
out_lse = tl.log(out_se) + max_lse
@@ -84,8 +87,8 @@ def merge_attn_states_kernel(
87
# NOTE(woosuk): Be careful with the numerical stability.
88
# We should compute the scale first, and then multiply it with the output.
89
# Do not multiply the output with tl.exp(p_lse) or tl.exp(s_lse) directly.
87
- p_scale = tl.exp(p_lse) / out_se
88
- s_scale = tl.exp(s_lse) / out_se
90
+ p_scale = p_se / out_se
91
+ s_scale = s_se / out_se
92
out = p_out * p_scale + s_out * s_scale
93
tl.store(output + token_idx * num_heads * HEAD_SIZE +
94
head_idx * HEAD_SIZE + head_arange,