Hack for flash-attn support for inference

ecmwf · Oct 15, 2024 · 54b1fd5 · 54b1fd5
1 parent f31f867
commit 54b1fd5
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/anemoi/models/layers/attention.py b/src/anemoi/models/layers/attention.py
@@ -89,9 +89,13 @@ def forward(
         dropout_p = self.dropout_p if self.training else 0.0
 
         if _FLASH_ATTENTION_AVAILABLE:
+            import torch
             query, key, value = (
                 einops.rearrange(t, "batch heads grid vars -> batch grid heads vars") for t in (query, key, value)
             )
+            query = query.to(torch.bfloat16, non_blocking=True).to(device="cuda")
+            key = key.to(torch.bfloat16, non_blocking=True).to(device="cuda")
+            value = value.to(torch.bfloat16, non_blocking=True).to(device="cuda")
             out = self.attention(query, key, value, causal=False, window_size=self.window_size, dropout_p=dropout_p)
             out = einops.rearrange(out, "batch grid heads vars -> batch heads grid vars")
         else: