Skip to content

[Kernel] Remove redundant Exp calculations #16123

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Apr 15, 2025
Merged

[Kernel] Remove redundant Exp calculations #16123

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
456e5d2
[Kernel] Remove redundant Exp calculations
DefTruth Apr 6, 2025
2aec245
[Kernel] Remove redundant Exp calculations
DefTruth Apr 6, 2025
5375375
Merge branch 'vllm-project:main' into remove-redundant-exp
DefTruth Apr 13, 2025
5714483
Merge branch 'vllm-project:main' into remove-redundant-exp
DefTruth Apr 14, 2025
d4b62f7
Merge branch 'vllm-project:main' into remove-redundant-exp
DefTruth Apr 14, 2025
34bfac9
Merge branch 'vllm-project:main' into remove-redundant-exp
DefTruth Apr 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions vllm/attention/ops/triton_merge_attn_states.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,10 @@ def merge_attn_states_kernel(
max_lse = tl.maximum(p_lse, s_lse)
p_lse = p_lse - max_lse
s_lse = s_lse - max_lse
out_se = (tl.exp(p_lse) + tl.exp(s_lse))
# Will reuse precomputed Exp values for scale factor computation.
p_se = tl.exp(p_lse)
s_se = tl.exp(s_lse)
out_se = (p_se + s_se)

if OUTPUT_LSE:
out_lse = tl.log(out_se) + max_lse
Expand All @@ -84,8 +87,8 @@ def merge_attn_states_kernel(
# NOTE(woosuk): Be careful with the numerical stability.
# We should compute the scale first, and then multiply it with the output.
# Do not multiply the output with tl.exp(p_lse) or tl.exp(s_lse) directly.
p_scale = tl.exp(p_lse) / out_se
s_scale = tl.exp(s_lse) / out_se
p_scale = p_se / out_se
s_scale = s_se / out_se
out = p_out * p_scale + s_out * s_scale
tl.store(output + token_idx * num_heads * HEAD_SIZE +
head_idx * HEAD_SIZE + head_arange,
Expand Down