[AMD] Don't request the unused softmax LSE in the AITER diffusion backend (#32926)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: YC Yen-Ching Tseng <yctseng@amd.com>
This commit is contained in:
co-authored by
Claude Opus 5
YC Yen-Ching Tseng
parent
13d2ee8180
commit
2d6c12e2fd
@@ -204,14 +204,14 @@ class AITerImpl(AttentionImpl):
|
||||
)
|
||||
|
||||
# BF16 path
|
||||
output, _ = aiter.flash_attn_func(
|
||||
output = aiter.flash_attn_func(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
dropout_p=self.dropout_p,
|
||||
causal=self.causal,
|
||||
return_attn_probs=False,
|
||||
return_lse=True,
|
||||
return_lse=False,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
Reference in New Issue
Block a user