FlashAttention-2: Attention ที่เร็วขึ้นด้วยการประมวลผลแบบขนานและการแบ่งงานที่ดีกว่า

• มีการปรากฏตัวของโมเดลภาษาที่มีคอนเท็กซ์ยาวขึ้น เช่น GPT-4(32k), MPT(65k), Calude(100k)
• การขยายความยาวคอนเท็กซ์ของทรานส์ฟอร์เมอร์ทำได้ยาก เพราะความต้องการด้านรันไทม์และหน่วยความจำเพิ่มขึ้นเป็นกำลังสี่
• FlashAttention ที่เปิดตัวเมื่อปีที่แล้วถูกใช้งานอย่างแพร่หลาย เพราะช่วยลดการใช้หน่วยความจำและเพิ่มความเร็วของ attention
• ตอนเปิดตัวนั้นเร็วกว่าเดิม 2-4 เท่าอยู่แล้ว แต่ยังมีช่องให้ปรับปรุงได้อีก เมื่อเทียบกับการคูณเมทริกซ์แบบปรับแต่งประสิทธิภาพ (GEMM) ก็ยังไม่เร็วพอ และในทางทฤษฎียังทำได้เพียง 25-40% ของค่า FLOPs/s สูงสุด (บน A100 GPU สูงสุด 124 TFLOPs/s)
• FlashAttention-2 เร็วกว่าเวอร์ชันก่อนหน้า 2 เท่า และให้ประสิทธิภาพสูงสุด 230 TFLOP/s บน A100 GPU
• ในการฝึกโมเดลภาษาแบบ GPT สามารถทำได้สูงสุดถึง 225 TFLOPS (ประสิทธิภาพการใช้ model FLOP 72%)
• ปรับอัลกอริทึมเพื่อลด non-matmul FLOPs
• ทำการขนานงานได้ดีขึ้น และเปลี่ยนวิธีแบ่งงานในแต่ละ thread block
• ขยายจำนวน Head Dimensions จาก 128 เป็น 256