Fix potential inconsistency with different attention mentods in hidream01

This commit is contained in:
Jaret Burkett 2026-05-13 09:08:11 -06:00
parent 4bfe944792
commit e03c6e4dc9
1 changed files with 10 additions and 0 deletions

View File

@ -170,6 +170,7 @@ def eager_attention_forward(
attention_mask: Optional[torch.Tensor],
scaling: float,
dropout: float = 0.0,
is_causal: bool = False,
**kwargs: Unpack[TransformersKwargs],
):
key_states = repeat_kv(key, module.num_key_value_groups)
@ -179,6 +180,15 @@ def eager_attention_forward(
if attention_mask is not None:
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
attn_weights = attn_weights + causal_mask
elif is_causal:
q_len = query.shape[-2]
k_len = key_states.shape[-2]
causal_mask = torch.ones(
q_len, k_len, dtype=torch.bool, device=query.device
).triu(diagonal=k_len - q_len + 1)
attn_weights = attn_weights.masked_fill(
causal_mask, torch.finfo(attn_weights.dtype).min
)
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
query.dtype