Upload modeling_flash_llama.py

modeling_flash_llama.py

@@ -369,6 +369,12 @@ class LlamaAttention(nn.Module):
             
             unpadded_kv, indices_k, cu_seqlens_k, max_seqlen_k = unpad_input(kv, attention_mask)
             unpadded_q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, attention_mask[:, -q.size(1):])
+            # cast kv and q to bf16 or fp16 if currently in float32
+            if unpadded_kv.dtype == torch.float32:
+                unpadded_kv = unpadded_kv.to(torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16)
+                unpadded_q = unpadded_q.to(torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16)
+
+            
             attn_outputs = flash_attn_varlen_kvpacked_func(
                 unpadded_q, unpadded_kv, cu_seqlens_q, cu_seqlens_k, 
                 max_seqlen_q, max_seqlen_k,
@@ -385,6 +391,10 @@ class LlamaAttention(nn.Module):
         else:
 
             # no padding tokens, more efficient
+            # cast to bf16 or fp16 if currently in float32
+            if kv.dtype == torch.float32:
+                kv = kv.to(torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16)
+                q = q.to(torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16)
             
             attn_outputs = flash_attn_kvpacked_func(
                 q, kv, dropout_p=0.0, softmax_scale=1.0/self.norm_factor, causal=(not has_layer_past), return_attn_probs=output_attentions)