FlashAttention-4 เปิดตัวอย่างน่าตื่นเต้น: ความเร็วของกลไกความสนใจบน GPU Blackwell เทียบเท่ากับการคูณเมทริกซ์ เพิ่มประสิทธิภาพสูงสุดถึง 2.7 เท่า

หลังจากพัฒนามาเป็นเวลาหนึ่งปี FlashAttention-4 ได้เปิดตัวอย่างเป็นทางการแล้ว

ในฐานะเทคโนโลยีการปรับปรุงประสิทธิภาพระดับล่างที่สำคัญในสาขา Deep Learning FlashAttention ได้รับการอัปเดตเวอร์ชันครั้งใหญ่ Tri Dao ผู้เขียนหลักและผู้ช่วยศาสตราจารย์จากมหาวิทยาลัย Princeton ระบุว่า บน GPU Blackwell ความเร็วในการดำเนินการของกลไกความสนใจ (Attention) ในขณะนี้เกือบจะเทียบเท่ากับการคูณเมทริกซ์แล้ว แม้ว่าจุดคอขวดของทั้งสองจะแตกต่างกันโดยสิ้นเชิง

ในปัจจุบัน ความเร็วของ Tensor Core นั้นเร็วมากจนจุดคอขวดของ Forward Pass ของ Attention ได้ย้ายไปอยู่ที่หน่วยคำนวณ Exponential (ฟังก์ชันเลขชี้กำลัง) ในขณะที่จุดคอขวดของ Backward Pass ของ Attention อยู่ที่แบนด์วิดท์ของ Shared Memory

อัลกอริทึมใหม่ได้รับการออกแบบใหม่เพื่อแก้ไขจุดคอขวดเหล่านี้ โดยนำกลไกการเพิ่มประสิทธิภาพหลายประการมาใช้ รวมถึงการใช้พหุนามเพื่อประมาณค่า Exponential อัลกอริทึม Online Softmax ใหม่เพื่อลดการดำเนินการ Rescaling ลง 90% และการใช้คำสั่ง 2CTA MMA เพื่อให้บล็อกเธรดสองบล็อกแชร์โอเปอแรนด์ร่วมกันเพื่อลดการไหลของข้อมูลใน Shared Memory เป็นต้น

• ที่อยู่บทความ: https://github.com/Dao-AILab/flash-attention/blob/main/assets/fa4_paper.pdf
• ลิงก์โค้ด: https://github.com/Dao-AILab/flash-attention

ต่อไป เราจะวิเคราะห์รายละเอียดทางเทคนิคของ FlashAttention-4 อย่างละเอียด

แนวโน้มฮาร์ดแวร์: การขยายแบบไม่สมมาตร

เป็นเวลานานมาแล้วที่กลไกความสนใจ (Attention) ซึ่งเป็นชั้นหลักของสถาปัตยกรรม Transformer ได้เป็นจุดคอขวดด้านประสิทธิภาพหลักสำหรับโมเดลภาษาขนาดใหญ่และแอปพลิเคชันที่มีคอนเท็กซ์ยาว ก่อนหน้านี้ FlashAttention-3 ได้เพิ่มประสิทธิภาพสำหรับสถาปัตยกรรม GPU Hopper (เช่น H100) ผ่านเทคโนโลยีเช่นการดำเนินการแบบอะซิงโครนัสและการเชี่ยวชาญเฉพาะ Warp

อย่างไรก็ตาม อุตสาหกรรม AI กำลังเปลี่ยนไปใช้ระบบสถาปัตยกรรม Blackwell (เช่น B200 และ GB200) อย่างรวดเร็ว Acelerator สมัยใหม่อย่าง GPU Blackwell สืบทอดแนวโน้มหนึ่ง: การขยายตัวของฮาร์ดแวร์แบบไม่สมมาตร (Asymmetric Scaling) ภายใต้แนวโน้มนี้ Throughput ของ Tensor Core จะเติบโตเร็วกว่าทรัพยากรฮาร์ดแวร์อื่นๆ มาก เช่น แบนด์วิดท์ของ Shared Memory หน่วยฟังก์ชันพิเศษ (Special Function Unit – SFU) สำหรับการคำนวณฟังก์ชัน Transcendental เช่น Exponential และ ALU จำนวนเต็มและจุดลอยตัวทั่วไป

ตัวอย่างเช่น จาก Hopper H100 ไปยัง Blackwell B200 Throughput ของ Tensor Core BF16 เพิ่มขึ้น 2.25 เท่า แต่จำนวน SFU และแบนด์วิดท์ Shared Memory ยังคงไม่เปลี่ยนแปลงโดยพื้นฐาน ความไม่สมดุลของการขยายตัวนี้ส่งผลกระทบอย่างลึกซึ้งต่อการเพิ่มประสิทธิภาพเคอร์เนลที่ซับซ้อนเช่น Attention

โดยเฉพาะอย่างยิ่ง แกนกลางของกลไก Attention ประกอบด้วยการดำเนินการ General Matrix Multiplication (GEMM) สองครั้ง โดยมี Softmax คั่นกลาง แต่ในทางปฏิบัติ Attention ยังเกี่ยวข้องกับงานเสริมจำนวนมาก เช่น การเคลื่อนย้ายข้อมูล การซิงโครไนซ์ การแปลงเลย์เอาต์ การดำเนินการระดับองค์ประกอบ การจัดตารางเวลา และการจัดการมาสก์ เป็นต้น

มุมมองดั้งเดิมเชื่อว่าประสิทธิภาพของ Attention ถูกกำหนดโดยความเร็วของ GEMM เท่านั้น อย่างไรก็ตาม การวิเคราะห์ “Speed and Feed” ของ B200 แสดงให้เห็นว่าจุดคอขวดหลักไม่ใช่ Tensor Core แต่คือ:
1. หน่วย SFU สำหรับการคำนวณ Exponential ของ Softmax ใน Forward Pass
2. การไหลของข้อมูล Shared Memory ที่ถูกจำกัดด้วยแบนด์วิดท์ Shared Memory ใน Backward Pass

เพื่อแก้ไขปัญหานี้ ทีมวิจัยได้เปิดตัว FlashAttention-4 ซึ่งเป็นการออกแบบร่วมกันของอัลกอริทึมและเคอร์เนล เป้าหมายหลักคือการบรรลุประสิทธิภาพสูงสุด 1605 TFLOPs/s บน B200 (BF16) (อัตราการใช้ประโยชน์ 71%) โดยการเพิ่มการทับซ้อนระหว่างการคูณเมทริกซ์กับทรัพยากรจุดคอขวดอื่นๆ ให้มากที่สุด ซึ่งเร็วกว่า cuDNN 9.13 ถึง 1.3 เท่า และเร็วกว่า Triton 2.7 เท่า

แนวคิดหลักของการออกแบบร่วมมีดังนี้:
* Pipeline ใหม่: ออกแบบซอฟต์แวร์ Pipeline ใหม่สำหรับ Forward และ Backward Pass โดยใช้ MMA แบบอะซิงโครนัสเต็มรูปแบบของ Blackwell และขนาดบล็อกที่ใหญ่ขึ้น เพื่อเพิ่มการทับซ้อนระหว่างการคำนวณ Tensor Core การคำนวณ Softmax และการดำเนินการหน่วยความจำให้มากที่สุด
* การเพิ่มประสิทธิภาพ Forward Pass: จำลองซอฟต์แวร์ของฟังก์ชัน Exponential บนหน่วย FMA ผ่านการประมาณค่าด้วยพหุนาม เพื่อเพิ่ม Throughput การคำนวณ Exponential พร้อมกันนี้ได้แนะนำ Conditional Softmax Rescaling เพื่อข้ามการดำเนินการ Rescaling ที่ไม่จำเป็น ซึ่งช่วยบรรเทาจุดคอขวดของ SFU
* การเพิ่มประสิทธิภาพ Backward Pass: ใช้ Tensor Memory (TMEM) เพื่อจัดเก็บผลลัพธ์ขั้นกลาง เพื่อบรรเทาความกดดันของการไหลของข้อมูล Shared Memory รวมกับโหมด 2-CTA MMA ใหม่ของ Blackwell เพื่อลดการเข้าถึง Shared Memory เพิ่มเติม และลดจำนวนการดำเนินการ Atomic Reduction ลงครึ่งหนึ่ง นอกจากนี้ยังรองรับโหมดการดำเนินการแบบ Deterministic เพื่อให้การฝึกซ้อมสามารถทำซ้ำได้
* การเพิ่มประสิทธิภาพการจัดตารางเวลา: แนะนำตัวจัดตารางเวลาแบบแบ่งบล็อกใหม่ เพื่อแก้ไขปัญหาการกระจายโหลดที่ไม่สมดุลที่เกิดจาก Causal Mask และลำดับความยาวแปรผัน

คุณสมบัติฮาร์ดแวร์ใหม่ของ Blackwell

• Tensor Memory (TMEM): บน B200 แต่ละ Streaming Multiprocessor (SM) มี TMEM ขนาด 256 KB เชื่อมต่อกับ Tensor Core โดยตรง ใช้สำหรับจัดเก็บผลลัพธ์ขั้นกลางที่ซิงโครไนซ์ในระดับ Warp
• Tensor Core รุ่นที่ 5 แบบอะซิงโครนัสเต็มรูปแบบ: คำสั่ง tcgen05.mma รองรับการดำเนินการแบบอะซิงโครนัส และจัดเก็บผลลัพธ์สะสมใน TMEM สำหรับ BF16 และ FP16 ขนาดบล็อก UMMA ใหญ่สุดที่ CTA เดียวสามารถใช้ได้คือ 128×256×16 ซึ่งใหญ่กว่าขนาดบล็อกอะตอมิก WGMMA ใหญ่สุดในสถาปัตยกรรม Hopper ประมาณ 2 เท่า UMMA ถูกเริ่มต้นโดยเธรดเดียว ซึ่งลดแรงกดดันต่อรีจิสเตอร์ ทำให้ง่ายต่อการใช้บล็อกที่ใหญ่กว่าและ Pipeline ที่ลึกขึ้นโดยไม่เกิดรีจิสเตอร์โอเวอร์โฟลว์ สิ่งนี้ยังทำให้การเชี่ยวชาญเฉพาะ Warp เป็นไปได้มากขึ้น โดย Warp บางส่วนรับผิดชอบในการเคลื่อนย้ายข้อมูล และ Warp อื่นๆ รับผิดชอบในการเริ่มต้น MMA เพื่อให้เกิดการทับซ้อนระหว่างการคำนวณและการเข้าถึงหน่วยความจำ tcgen05.mma ยังสามารถอ่านโอเปอแรนด์ A จาก TMEM โดยตรงได้
• 2-CTA MMA: Blackwell รองรับให้ CTA คู่หนึ่งในคลัสเตอร์เดียวกันดำเนินการ UMMA ร่วมกัน และครอบคลุม TMEM ของทั้งสอง CTA โดยเธรดหนึ่งใน Leader CTA จะเริ่มต้น MMA แต่ทั้งสอง CTA ต้องยังคงทำงานอยู่ระหว่างการดำเนินการ ด้วยการแบ่งมิติ M และ N ระหว่าง CTA คู่นี้ สามารถขยายขนาดบล็อกของ MMA เป็น 256×256×16 ซึ่งช่วยลดการถ่ายโอนข้อมูลที่ซ้ำซ้อนและลดการใช้ทรัพยากรของแต่ละ CTA ภายในเคอร์เนลเดียว ขนาดของกลุ่ม CTA (1 หรือ 2) ต้องสอดคล้องกันระหว่างการดำเนินการ TMEM และการดำเนินการ Tensor Core

ภาษาและเฟรมเวิร์กการเขียนโปรแกรม: CuTe-DSL

FlashAttention-4 ถูกนำไปใช้ทั้งหมดโดยใช้ CuTe-DSL ซึ่งเป็นภาษาเฉพาะโดเมน (Domain-Specific Language – DSL) สำหรับเคอร์เนลที่มาจาก CUTLASS โค้ดเคอร์เนลถูกเขียนด้วย Python จากนั้น DSL จะลดระดับลงเป็น PTX และคอมไพล์เป็นโค้ดเครื่อง GPU โดย CUDA Toolchain

โมเดลการเขียนโปรแกรมนี้สอดคล้องกับ CuTe / CUTLASS ในระดับนามธรรม ในขณะเดียวกันก็ให้อินเทอร์เฟซควบคุมระดับต่ำในระดับ PTX เมื่อเทียบกับการใช้เทมเพลต C++ วิธีนี้สามารถลดเวลาในการคอมไพล์ลงได้ประมาณ 20–30 เท่า Tri Dao รู้สึกตื่นเต้นกับเรื่องนี้ เพราะทำให้กระบวนการติดตั้งและ “คอมไพล์” ใช้เวลาเพียงไม่กี่วินาที แทนที่จะเป็นหลายนาทีหรือหลายชั่วโมงเหมือนในอดีต

การทดสอบประสิทธิภาพมาตรฐานของ Attention

ทีมวิจัยได้แสดงผลลัพธ์ประสิทธิภาพของ FlashAttention-4 บน B200 (BF16) และเปรียบเทียบกับ FlashAttention-2 รวมถึงการนำไปใช้ของ Triton, Gluon และ cuDNN ผลลัพธ์แสดงให้เห็นว่า:

• ประสิทธิภาพ Forward Pass: บน GPU Blackwell Forward Pass ของ FlashAttention-4 เร็วกว่า cuDNN 9.13 1.1–1.3 เท่า และเร็วกว่าการนำไปใช้ของ Triton 2.1–2.7 เท่า
• ประสิทธิภาพ Backward Pass: เมื่อประมวลผลลำดับยาว ประสิทธิภาพ Backward Pass ของ FlashAttention-4 ยังคงเหนือกว่าวิธีมาตรฐานอื่นๆ อย่างต่อเนื่อง

การเปิดตัว FlashAttention-4 ได้รับความสนใจอย่างกว้างขวาง ทีม PyTorch ประกาศว่าฟีเจอร์ FlexAttention ของพวกเขารองรับแบ็กเอนด์ FlashAttention-4 แล้ว

FlexAttention ช่วยให้นักวิจัยสามารถสร้างต้นแบบรูปแบบต่างๆ ของกลไก Attention ที่กำหนดเองได้อย่างรวดเร็วมาเป็นเวลานาน ถูกนำไปใช้โดยโค้ดเบสหลายพันแห่งและได้รับการอ้างอิงในบทความหลายสิบบทความ แต่ผู้ใช้มักเผชิญกับจุดคอขวดด้านประสิทธิภาพ ด้วยการเปิดตัว FlashAttention-4 ทีม PyTorch ได้รวมแบ็กเอนด์นี้เข้ากับ FlexAttention บน GPU Hopper และ Blackwell ตอนนี้ PyTorch สามารถสร้างโค้ดแก้ไข score/mask ของ CuTeDSL โดยอัตโนมัติ และสร้างอินสแตนซ์ FlashAttention-4 สำหรับรูปแบบ Attention ที่กำหนดเองผ่านการคอมไพล์ JIT

ผลการทดสอบแสดงให้เห็นว่า ภายใต้เวิร์กโหลดที่ถูกจำกัดด้วยการคำนวณ เมื่อเทียบกับการนำไปใช้ของ Triton การเปลี่ยนแปลงนี้ยังคงนำมาซึ่งการเพิ่มประสิทธิภาพ 1.2 ถึง 3.2 เท่า ทำให้นักวิจัยไม่ต้องเลือกระหว่าง “ความยืดหยุ่น” กับ “ประสิทธิภาพสูง”

มีความคิดเห็นชี้ให้เห็นว่า FlashAttention-4 เป็นจุดหมายสำคัญ บนสถาปัตยกรรม Blackwell ความเร็วในการคำนวณ Attention ใกล้เคียงกับการคูณเมทริกซ์แล้ว ซึ่งหมายความว่าจุดคอขวดด้านการคำนวณจะถูกย้ายไปยังแบนด์วิดท์หน่วยความจำและการสื่อสารโดยสมบูรณ์ ประสิทธิภาพ Attention ที่ประมาณ 1600 TFLOPs ของมัน เมื่อเทียบกับ FlashAttention-3 แล้วเพิ่มขึ้น 2–3 เท่า ซึ่งจะส่งผลดีโดยตรงต่อโมเดลภาษาขนาดใหญ่ล้ำสมัยทั้งหมด เพราะหมายถึงหน้าต่างคอนเท็กซ์ที่มีประสิทธิภาพยาวนานขึ้น ค่าใช้จ่ายในการอนุมานที่ต่ำลง และความสามารถในการอนุมานแบบขยายขนาดได้ดีขึ้น

⚠️ หมายเหตุ: เนื้อหาได้รับการแปลโดย AI และตรวจสอบโดยมนุษย์ หากมีข้อผิดพลาดโปรดแจ้ง