Cleanup FlashAttention 2/3 a bit

flue-flash-attn-v2/src/lib.rs

@@ -143,10 +143,8 @@ impl FlashAttn {
         let seqlen_k_rounded = round_multiple(seqlen_k, 128);
 
         let elem_count = out_shape.elem_count();
-        let dst = unsafe { dev.alloc::<T>(elem_count) }.w()?;
-        let softmax_lse = dev
-            .alloc_zeros::<f32>(b_sz * 128 * num_heads * seqlen_q)
-            .w()?;
+        let dst = unsafe { dev.alloc::<T>(elem_count) }?;
+        let softmax_lse = dev.alloc_zeros::<f32>(b_sz * 128 * num_heads * seqlen_q)?;
 
         let is_bf16 = if is_bf16 { 1 } else { 0 };
 
@@ -610,10 +608,8 @@ impl FlashAttnVarLen {
         let seqlen_k_rounded = round_multiple(self.max_seqlen_k, 128);
 
         let elem_count = out_shape.elem_count();
-        let dst = unsafe { dev.alloc::<T>(elem_count) }.w()?;
-        let softmax_lse = dev
-            .alloc_zeros::<f32>(batch_size * num_heads * self.max_seqlen_q)
-            .w()?;
+        let dst = unsafe { dev.alloc::<T>(elem_count) }?;
+        let softmax_lse = dev.alloc_zeros::<f32>(batch_size * num_heads * self.max_seqlen_q)?;
 
         let is_bf16 = if is_bf16 { 1 } else { 0 };