A better fix for #3 to support new flash-attn versions
#38
+3
−3
flash-attn versions
#38