FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao; Daniel Y. Fu; Stefano Ermon; Atri Rudra; Christopher Ré

งานวิจัย · arXiv 2205.14135

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré

เผยแพร่ 2022-05-27cs.LG arXiv PDF

อ่านพร้อมคำแปลและ AI

สรุปสามบรรทัด

🎯 Standard exact attention ต้องสร้างเมทริกซ์ attention ขนาด $O (N^{2})$ ขึ้นจริงใน GPU HBM ที่มีความเร็วต่ำ ทำให้ติดคอขวดที่หน่วยความจำ (memory-bound) และขยายขนาดเพื่อรองรับบริบทขนาดยาวได้ยาก
🔧 FlashAttention คำนวณ exact attention ได้ภายใน CUDA kernel เดียวที่ถูกรวมเข้าด้วยกัน (fused) โดยการแบ่งอินพุตเป็นบล็อก (tiling) ลงบน on-chip SRAM ที่รวดเร็ว อัปเดต softmax ทีละส่วน และคำนวณ attention ซ้ำบนชิปในรอบ backward pass
📈 FlashAttention เร่งความเร็วการฝึกแบบ end-to-end ได้สูงสุดถึง 3 เท่าบน GPT-2, ฝึก BERT-large ได้เร็วกว่าสถิติของ MLPerf 1.1 ถึง 15% และขยายขนาดบริบทได้ถึง 64K โทเค็น

บทคัดย่อ

Transformers ทำงานช้าและใช้หน่วยความจำมหาศาลกับลำดับข้อมูลขนาดยาว เนื่องจากความซับซ้อนทั้งด้านเวลาและหน่วยความจำของ self-attention นั้นแปรผันตรงตามกำลังสองของความยาวลำดับ วิธี approximate attention พยายามแก้ปัญหานี้โดยยอมแลกคุณภาพของโมเดลเพื่อลดความซับซ้อนในการคำนวณ แต่บ่อยครั้งกลับไม่สามารถเร่งความเร็วตามเวลาจริง (wall-clock speedup) ได้เลย เรามองว่าหลักการที่ขาดหายไปคือการทำให้ขั้นตอนวิธีของ attention มีการตระหนักรู้เรื่อง IO (IO-aware) โดยคำนึงถึงการอ่านและการเขียนระหว่างระดับชั้นหน่วยความจำของ GPU เราจึงนำเสนอ FlashAttention ซึ่งเป็นขั้นตอนวิธีคำนวณ exact attention ที่เป็น IO-aware โดยใช้เทคนิค tiling เพื่อลดจำนวนครั้งของการอ่าน/เขียนหน่วยความจำระหว่าง high bandwidth memory (HBM) ของ GPU กับ on-chip SRAM ของ GPU เราวิเคราะห์ความซับซ้อนของ IO ใน FlashAttention และแสดงให้เห็นว่ามันเข้าถึง HBM น้อยกว่า standard attention อย่างมาก และมีประสิทธิภาพเหมาะสมที่สุด (optimal) สำหรับขนาด SRAM ในช่วงค่าต่างๆ นอกจากนี้ เรายังได้ขยาย FlashAttention ไปสู่ block-sparse attention ทำให้ได้ขั้นตอนวิธี approximate attention ที่เร็วกว่าวิธี approximate attention อื่นๆ ที่มีอยู่ทั้งหมด FlashAttention ช่วยฝึก Transformers ได้เร็วกว่าแบบจำลองพื้นฐานเดิมอย่างชัดเจน: เร่งความเร็ว wall-clock แบบครบวงจร (end-to-end) ขึ้น 15% บน BERT-large (ความยาวลำดับ 512) เมื่อเทียบกับสถิติความเร็วการฝึกของ MLPerf 1.1, เร็วขึ้น 3 $\times$ บน GPT-2 (ความยาวลำดับ 1K) และเร็วขึ้น 2.4 $\times$ บน Long-Range Arena (ความยาวลำดับ 1K-4K) นอกจากนี้ FlashAttention และ block-sparse FlashAttention ยังเปิดโอกาสให้ Transformers ใช้บริบทที่ยาวขึ้น ส่งผลให้โมเดลมีคุณภาพสูงขึ้น (ได้ค่า perplexity บน GPT-2 ดีขึ้น 0.7 และเพิ่มความแม่นยำ 6.4 จุดในการจำแนกเอกสารขนาดยาว) อีกทั้งยังปลดล็อกขีดความสามารถใหม่อย่างที่ไม่เคยมีมาก่อน โดยเป็น Transformers กลุ่มแรกที่ทำผลงานได้ดีกว่าการสุ่มเดาในโจทย์ท้าทายอย่าง Path-X (ความยาวลำดับ 16K, ความแม่นยำ 61.4%) และ Path-256 (ความยาวลำดับ 64K, ความแม่นยำ 63.1%)

บทคัดย่อต้นฉบับ (ภาษาอังกฤษ)

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3 $\times$ speedup on GPT-2 (seq. length 1K), and 2.4 $\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).

คู่มือการอ่าน

ปัญหาและแรงจูงใจ

Transformers มีความซับซ้อนทั้งด้านเวลาและหน่วยความจำแปรผันกำลังสองตามความยาวลำดับ $N$ แม้วิธี approximate attention จะช่วยลดจำนวน FLOP ในเชิงทฤษฎีได้ แต่บ่อยครั้งกลับไม่สามารถเร่งความเร็วตามเวลาจริง (wall-clock speedup) ได้เลย เนื่องจากคอขวดหลักของการคำนวณบน GPU ยุคใหม่เกิดจากการถ่ายโอนข้อมูลในหน่วยความจำ (IO) มากกว่าการดำเนินการทางคณิตศาสตร์

ในการคำนวณ standard attention การคำนวณ $S = Q K^{⊤}$ , $P = softmax (S)$ และ $O = PV$ จะต้องเขียนเมทริกซ์ตัวกลางขนาด $N \times N$ ลงใน High Bandwidth Memory (HBM) ที่มีความเร็วต่ำ และอ่านกลับขึ้นมาใหม่หลายครั้ง ซึ่งก่อให้เกิดการเข้าถึง HBM ถึง $Θ (N d + N^{2})$ ครั้ง และใช้พื้นที่หน่วยความจำ $O (N^{2})$ จึงกลายเป็นข้อจำกัดหลักของหน้าต่างบริบท (context window) ในทางปฏิบัติ

แนวคิดสำคัญ

IO-Awareness: คำนึงถึงลำดับขั้นหน่วยความจำของ GPU อย่างชัดเจน ได้แก่ on-chip SRAM ที่มีขนาดเล็กและรวดเร็วมาก (เช่น 192 KB ต่อ SM ด้วยแบนด์วิดท์ ~19 TB/s บน A100) เทียบกับ HBM ที่มีขนาดใหญ่แต่ช้ากว่า (40–80 GB ด้วยแบนด์วิดท์ 1.5–2.0 TB/s)
Tiling (การคำนวณ Softmax แบบแบ่งบล็อก): แบ่งอินพุต $Q, K, V$ ออกเป็นบล็อกๆ ที่มีขนาดพอดีกับ SRAM การติดตามสถิติค่าสูงสุดประจำแถวและตัวปรับบรรทัดฐานของ softmax แบบต่อเนื่อง ช่วยให้สามารถคำนวณ softmax reduction สะสมข้ามบล็อกได้โดยไม่ต้องสร้างเมทริกซ์ขนาด $N \times N$ ทั้งหมดลงบน HBM
Recomputation in Backward Pass: หลีกเลี่ยงการจัดเก็บเมทริกซ์ attention ตัวกลางขนาด $N \times N$ สำหรับใช้ในกระบวนการ backpropagation แต่จะจัดเก็บเฉพาะเอาต์พุต $O$ และสถิติของ softmax $(ℓ, m)$ เอาไว้ แล้วทำการคำนวณ $S$ และ $P$ ซ้ำทีละบล็อกบน SRAM ภายในชิประหว่างรอบ backward pass แทน
Kernel Fusion: รวมการคูณเมทริกซ์, masking, softmax และ dropout เข้าเป็น CUDA kernel เดียว เพื่อตัดวงจรการอ่านและเขียนกลับไปยัง HBM ซ้ำซ้อน

ระเบียบวิธี (พร้อมสมการสำคัญ)

เมื่อกำหนด $Q, K, V \in R^{N \times d}$ FlashAttention จะกำหนดขนาดบล็อก $B_{c} = ⌈ M / (4 d)⌉$ และ $B_{r} = min (⌈ M / (4 d)⌉, d)$ โดยที่ $M$ คือความจุของ SRAM

การติดตาม Softmax แบบสะสม (Incremental Softmax Tracking)

สำหรับเวกเตอร์ที่ถูกแยกย่อยออกเป็นเวกเตอร์ย่อย $x^{(1)}$ และ $x^{(2)}$ สถิติสะสมจะอัปเดตผ่าน: $m (x) = max (m (x^{(1)}), m (x^{(2)})),$ $ℓ (x) = e^{m (x^{(1)}) - m (x)} ℓ (x^{(1)}) + e^{m (x^{(2)}) - m (x)} ℓ (x^{(2)}) .$

สำหรับบล็อก $j$ ของ keys/values และบล็อก $i$ ของ queries:

คำนวณคะแนนที่ยังไม่ได้ปรับบรรทัดฐานบน SRAM: $S_{ij} = Q_{i} K_{j}^{⊤} \in R^{B_{r} \times B_{c}}$
คำนวณสถิติเฉพาะที่ของบล็อก: $\tilde{m}_{ij} = rowmax (S_{ij})$ , $\tilde{P}_{ij} = exp (S_{ij} - \tilde{m}_{ij})$ และ $\tilde{ℓ}_{ij} = rowsum (\tilde{P}_{ij})$
อัปเดตตัวปรับบรรทัดฐานสะสมระดับโกลบอล: $m_{i}^{new} = max (m_{i}, \tilde{m}_{ij}), ℓ_{i}^{new} = e^{m_{i} - m_{i}^{new}} ℓ_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{ℓ}_{ij} .$
อัปเดตบล็อกเอาต์พุตสะสม $O_{i} \in R^{B_{r} \times d}$ : $O_{i} \leftarrow diag (ℓ_{i}^{new})^{- 1} (diag (ℓ_{i}) e^{m_{i} - m_{i}^{new}} O_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{P}_{ij} V_{j}) .$

การวิเคราะห์ความซับซ้อน (Complexity Analysis)

ความซับซ้อนด้าน IO: Standard attention ต้องการการเข้าถึง HBM ระดับ $Θ (N d + N^{2})$ ในขณะที่ FlashAttention ต้องการการเข้าถึง HBM เพียง $Θ (N^{2} d^{2} M^{- 1})$ ผู้เขียนพิสูจน์แล้วว่าไม่มีขั้นตอนวิธี exact attention ใดที่จะเข้าถึง HBM เชิงเส้นกำกับได้ดีกว่า $o (N^{2} d^{2} M^{- 1})$ สำหรับทุกค่า $M \in [d, N d]$
FLOPs และหน่วยความจำ: FlashAttention มีความซับซ้อนในการคำนวณ $O (N^{2} d)$ FLOPs (มี FLOPs เพิ่มขึ้นเล็กน้อยในรอบ backward pass จากการคำนวณซ้ำ) และต้องการหน่วยความจำเสริมเพียง $O (N)$ นอกเหนือจากอินพุตและเอาต์พุต

ส่วนขยาย: Block-Sparse FlashAttention

การจำกัดขอบเขตการคำนวณเฉพาะบล็อกที่ไม่เป็นศูนย์ตามแบบแผนแมสก์ความเบาบางที่กำหนดไว้ $M \in {0, 1}^{N / B_{r} \times N / B_{c}}$ จะให้ความซับซ้อนด้าน IO อยู่ที่ $Θ (N d + N^{2} d^{2} M^{- 1} s)$ โดยที่ $s$ คือสัดส่วนของบล็อกที่ไม่เป็นศูนย์

ผลลัพธ์ (พร้อมตัวเลขที่สำคัญที่สุด)

การเร่งความเร็วในการฝึก:
- BERT-large (ความยาวลำดับ 512): เร็วกว่าสถิติความเร็วการฝึกของ MLPerf 1.1 ถึง 15% (17.4 นาที เทียบกับ 20.0 นาที บน $8 \times$ A100s)
- GPT-2 (ความยาวลำดับ 1K): เร็วกว่า HuggingFace สูงสุด 3.5 เท่า และเร็วกว่า Megatron-LM สูงสุด 1.8 เท่า
- Long-Range Arena (LRA): ทำงานเร็วกว่า standard attention 2.4 เท่า และ Block-sparse FlashAttention เร็วขึ้น 2.8 เท่า
ขีดความสามารถด้านบริบทที่ยาวขึ้น:
- GPT-2 พร้อมบริบท 4K: ฝึกได้เร็วกว่า Megatron-LM ที่บริบท 1K ขณะเดียวกันก็ปรับปรุง validation perplexity ดีขึ้น 0.7 จุด (17.5 เทียบกับ 18.2)
- การจำแนกเอกสารขนาดยาว: การขยายบริบทจาก 512 เป็น 16K ช่วยเพิ่ม micro- $F_{1}$ บน MIMIC-III ได้ 4.3 จุด (52.8 เป็น 57.1) และเพิ่มขึ้นสูงสุด 8.5 จุดบน ECtHR (72.2 เป็น 80.7 ที่ความยาว 8K)
- Path-X และ Path-256: FlashAttention เป็นแบบจำลอง Transformer ตัวแรกที่สามารถผ่านโจทย์ Path-X (ความยาวลำดับ 16K) ด้วยความแม่นยำ 61.4% และ Block-sparse FlashAttention ทำความแม่นยำได้ 63.1% บน Path-256 (ความยาวลำดับ 64K)

ข้อจำกัดและประเด็นที่ยังเปิดกว้าง

ภาระงานด้านวิศวกรรม CUDA: การพัฒนา IO-aware kernel ต้องอาศัยการเขียนโค้ด CUDA ระดับล่างแบบกำหนดเอง ไม่สามารถใช้เพียงนามธรรมระดับสูงของ PyTorch หรือคอมไพเลอร์ทั่วไป ส่งผลให้การย้ายข้ามสถาปัตยกรรมฮาร์ดแวร์ทำได้ยาก
Attention บน Multi-GPU: แผนการจัดการ IO ที่เหมาะสมที่สุดบนโหนดเดี่ยว ไม่สามารถนำไปขยายผลโดยตรงกับ attention แบบกระจายตัวบนหลาย GPU ที่เชื่อมต่อผ่านโครงข่ายระดับสูง (NVLink, InfiniBand)
ขอบเขตล่างของ Exact Attention: แม้จะมีประสิทธิภาพเหมาะสมที่สุดสำหรับค่า $M$ ทั่วไปแล้ว แต่การพิสูจน์ขอบเขตล่างของความซับซ้อนแบบแบ่งพารามิเตอร์อย่างละเอียดในรูปของ $M$ ยังคงเป็นคำถามที่เปิดกว้างอยู่

คำศัพท์สำคัญ

IO-awareness: หลักการทางขั้นตอนวิธีในการออกแบบการคำนวณให้คำนึงถึงการอ่านและการเขียนระหว่างระดับชั้นต่างๆ ของหน่วยความจำฮาร์ดแวร์ (เช่น on-chip SRAM ที่รวดเร็ว เทียบกับ HBM ที่ช้ากว่า)
High Bandwidth Memory (HBM): หน่วยความจำหลักขนาดใหญ่ของ GPU ซึ่งมีความจุรวมสูงมากแต่มีแบนด์วิดท์ต่ำกว่าหน่วยความจำแคชบนชิปอย่างมีนัยสำคัญ
SRAM: หน่วยความจำ scratchpad บนชิปที่มีความเร็วสูงใน GPU ยุคใหม่ ซึ่งติดตั้งอยู่ภายในแต่ละ Streaming Multiprocessor (SM)
Tiling: การปรับโครงสร้างกระบวนการคำนวณโดยแบ่งเมทริกซ์อินพุตออกเป็นบล็อกย่อยที่มีขนาดพอดีกับ SRAM แล้วนำผลลัพธ์มาคำนวณสะสมร่วมกันข้ามรอบการทำงาน
Selective gradient checkpointing: เทคนิคที่หลีกเลี่ยงการบันทึกเมทริกซ์ attention ตัวกลางขนาดใหญ่ในระหว่าง forward pass โดยเลือกที่จะคำนวณซ้ำบนชิประหว่าง backward pass โดยอาศัยค่าสถิติ softmax ที่บันทึกไว้แทน
Kernel fusion: การรวมการคำนวณหลายคำสั่งที่แยกจากกัน (เช่น matrix multiply, masking, softmax, dropout) ให้ทำงานเสร็จสิ้นภายใน GPU kernel เดียว เพื่อลดการอ่าน/เขียนข้อมูลกลับไปยัง HBM ซ้ำซ้อน
Block-sparse FlashAttention: ส่วนขยายของ FlashAttention ที่เลือกคำนวณเฉพาะบล็อกที่ไม่เป็นศูนย์ตามแมสก์ความเบาบางแบบบล็อก ซึ่งช่วยลดความซับซ้อนด้าน IO ลงตามสัดส่วนความเบาบางนั้น
Path-X / Path-256: ชุดทดสอบการให้เหตุผลเชิงภาพสำหรับลำดับขนาดยาวจากเบঞ্চมาร์ก Long-Range Arena (ความยาวลำดับ 16K และ 64K) ซึ่งในอดีตไม่มี Transformers รุ่นใดทำความแม่นยำได้สูงกว่าการสุ่มเดา

คำถามที่พบบ่อย

FlashAttention เป็นการประมาณค่าเมทริกซ์ attention ใช่หรือไม่?

ไม่ใช่ FlashAttention เป็นการคำนวณ exact self-attention อย่างแท้จริง โดยให้ผลลัพธ์ทางคณิตศาสตร์ที่เหมือนกับการคำนวณ standard attention ทุกประการ จนถึงระดับความแม่นยำของทศนิยมแบบ floating-point

เหตุใด FlashAttention จึงทำงานได้เร็วกว่า ทั้งที่มีการคำนวณ FLOPs ในรอบ backward pass มากกว่า?

เพราะ standard attention ติดคอขวดที่หน่วยความจำ (memory-bound) การเคลื่อนย้ายเมทริกซ์ขนาด $N \times N$ ไปมาระหว่าง HBM และ SRAM ใช้เวลาตามจริง (wall-clock time) นานกว่าการสั่งคำนวณทางคณิตศาสตร์เพิ่มเติมบน SRAM ที่รวดเร็วมาก

FlashAttention ต้องบันทึกหน่วยความจำเสริมส่วนใดไว้สำหรับรอบ backward pass บ้าง?

FlashAttention จะเก็บเฉพาะเอาต์พุต $O$ จากรอบ forward pass, ค่าแฟกเตอร์ปรับบรรทัดฐานของ softmax ได้แก่ $ℓ$ และ $m$ (ซึ่งใช้หน่วยความจำเพิ่มเพียง $O (N)$ ) และสถานะของตัวสร้างเลขสุ่มเทียมเพื่อจำลองแมสก์ dropout ขึ้นใหม่บนชิป

ความซับซ้อนด้าน IO ของ FlashAttention เปรียบเทียบกับ standard attention แล้วเป็นอย่างไร?

Standard attention ต้องการการเข้าถึง HBM ที่ $Θ (N d + N^{2})$ ครั้ง ในขณะที่ FlashAttention ต้องการการเข้าถึงเพียง $Θ (N^{2} d^{2} M^{- 1})$ ครั้ง โดยที่ $M$ คือขนาดความจุของ SRAM และ $d$ คือมิติของ head

คำอธิบายในหน้านี้เขียนโดย AI จากงานวิจัยและอาจมีข้อผิดพลาด โปรดอ่านต้นฉบับเพื่อดูรายละเอียด

อ่านพร้อมคำแปลและ AI