[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
|
# BF16 path
|
||||||
output, _ = aiter.flash_attn_func(
|
output = aiter.flash_attn_func(
|
||||||
query,
|
query,
|
||||||
key,
|
key,
|
||||||
value,
|
value,
|
||||||
dropout_p=self.dropout_p,
|
dropout_p=self.dropout_p,
|
||||||
causal=self.causal,
|
causal=self.causal,
|
||||||
return_attn_probs=False,
|
return_attn_probs=False,
|
||||||
return_lse=True,
|
return_lse=False,
|
||||||
)
|
)
|
||||||
return output
|
return output
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user