This commit is contained in:
Yina Chen 2024-11-14 04:03:36 +02:00 committed by GitHub
parent 00fce5c940
commit 59b01fa7d2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -934,9 +934,9 @@ class PrefillRunner:
hidden_states = F.pad(hidden_states.to(torch.float16), (0, 0, 0, pad_len), value=0.0) hidden_states = F.pad(hidden_states.to(torch.float16), (0, 0, 0, pad_len), value=0.0)
position_ids = F.pad(position_ids, (0, pad_len), value=0) position_ids = F.pad(position_ids, (0, pad_len), value=0)
attention_mask = F.pad( attention_mask = F.pad(
attention_mask.to(torch.int64), attention_mask.to(torch.float16),
(0, pad_len, 0, pad_len), (0, pad_len, 0, pad_len),
value=torch.iinfo(torch.int64).min, value=torch.finfo(torch.float16).min,
) )
args = (hidden_states, position_ids, attention_mask, past_key_value) args = (hidden_states, position_ids, attention_mask, past_key_value)