Upload modeling_flash_llama.py

modeling_flash_llama.py

@@ -361,47 +361,14 @@ class LlamaAttention(nn.Module):
 
         past_key_value = (past_kv, past_len+q.size(1)) if use_cache else None
 
-        if is_padded_inputs:
-
-            # varlen, ignore padding tokens, efficient for large batch with many paddings
-
-            assert attention_mask is not None
-            
-            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,
-                dropout_p=0.0, softmax_scale=1.0/self.norm_factor, 
-                causal=(not has_layer_past), return_attn_probs=output_attentions
-            )
-
-            attn_output = attn_outputs[0] if output_attentions else attn_outputs
-            attn_output = pad_input(
-                attn_output, indices_q, bsz, max_seqlen_q
-            ).reshape(bsz, q_len, h_size)
-            attn_weights = attn_outputs[2] if output_attentions else None
-            
-        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)
-
-            attn_output = attn_outputs[0] if output_attentions else attn_outputs
-            attn_output = attn_output.reshape(bsz, q_len, h_size)
-            attn_weights = attn_outputs[2] if output_attentions else None
+        # no padding tokens, more efficient
+        attn_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+        attn_outputs = flash_attn_kvpacked_func(
+            q.type(attn_dtype), kv.type(attn_dtype), dropout_p=0.0, softmax_scale=1.0/self.norm_factor, causal=(not has_layer_past), return_attn_probs=output_attentions)
+
+        attn_output = attn_outputs[0] if output_attentions else attn_outputs
+        attn_output = attn_output.reshape(bsz, q_len, h_size)
+        attn_weights = attn_outputs[2] if output_attentions else None
 
         if self.config.pretraining_tp > 1:
             attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)