Fix CPU Fallback for NewAttention with xformers BlockDiagonalMask

escription:
This PR addresses two critical issues when running the model on CPU:

1. Memory Efficient Attention CPU Fallback:
- Added CPU device detection in NewAttention initialization
- Automatically disables xformers memory efficient attention when running on CPU
- Prevents NotImplementedError from xformers which only supports CUDA devices
- Added proper handling of xformers BlockDiagonalMask in standard attention
- Materializes BlockDiagonalMask to tensor before addition with attention scores

The fix ensures smooth fallback to standard attention mechanism when running on CPU while maintaining compatibility with xformers mask types.

Related issue:
NotImplementedError: No operator found for `memory_efficient_attention_forward` with inputs:
query : shape=(1, 5471, 16, 64) (torch.float32)
key : shape=(1, 5471, 16, 64) (torch.float32)
value : shape=(1, 5471, 16, 64) (torch.float32)
attn_bias : <class 'xformers.ops.fmha.attn_bias.BlockDiagonalMask'>
p : 0.0
`fa2F@v2.5.7-pt` is not supported because:
device=cpu (supported: {'cuda'})
dtype=torch.float32 (supported: {torch.float16, torch.bfloat16})
`cutlassF-pt` is not supported because:
device=cpu (supported: {'cuda'})

modeling.py

@@ -445,6 +445,10 @@ class NewAttention(nn.Module):
 
         if use_memory_efficient_attention is None:
             use_memory_efficient_attention = self.config.use_memory_efficient_attention
+
+        if not torch.cuda.is_available() or (hasattr(config, 'device') and config.device == 'cpu'):
+            use_memory_efficient_attention = False
+            
         self.use_memory_efficient_attention = use_memory_efficient_attention
         self.memory_efficient_attention = None if xops is None else xops.memory_efficient_attention
         if self.use_memory_efficient_attention:
@@ -489,6 +493,9 @@ class NewAttention(nn.Module):
             key_states = pad_input(key_states.squeeze(), *padding_inputs)
             value_states = pad_input(value_states.squeeze(), *padding_inputs)
 
+        if self.use_memory_efficient_attention and not hidden_states.is_cuda:
+            self.use_memory_efficient_attention = False
+        
         if self.use_memory_efficient_attention:
             assert self.memory_efficient_attention is not None, "xformers is not loaded"
             assert output_attentions is False, "memory_efficient_attention do not output attentions"
@@ -534,6 +541,12 @@ class NewAttention(nn.Module):
 
         attention_scores = attention_scores / math.sqrt(self.attention_head_size)
         if attention_bias is not None:
+           if hasattr(attention_bias, 'materialize'):
+                # If it's a BlockDiagonalMask, materialize it to a tensor
+                attention_bias = attention_bias.materialize(
+                    (attention_scores.shape[0], attention_scores.shape[1], 
+                    attention_scores.shape[2], attention_scores.shape[3])
+                )
             # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
             attention_scores = attention_scores + attention_bias