บทความแนะนำการวิจัยชื่อ CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs โดยมีเป้าหมายหลักในการปรับปรุงประสิทธิภาพการฝึกโมเดล Transformer โดยเฉพาะการแก้ไขการดำเนินการที่ดูเหมือนกระจัดกระจายแต่สะสมเวลาในการประมวลผลอย่างมากซึ่งเป็นประเภทที่ใช้หน่วยความจำสูง

ผู้เขียนบทความ แหล่งที่มา: Machine Learning News

วันที่ 22 พฤษภาคม ทรี ดาโอ ได้แชร์ทวีตของฮาน กัว่บนโซเชียลมีเดีย เขายังเขียนเพิ่มเติมว่า: “หลังจากเขียนสมการใหม่บางส่วน พบว่าทุกอย่างใน Transformer ล้วนเป็นชุดของ GEMM + epilogue (การคูณเมทริกซ์บวกส่วนท้าย) โดยมีพริมีทีฟที่ได้รับการปรับแต่งแล้ว LLM (และผู้เริ่มต้น) สามารถเขียนเคอร์เนลความเร็วแสงสำหรับการดำเนินการทั้งหมดของ Transformer!”

Tri Dao เป็นหนึ่งในผู้เขียนหลักของซีรีส์ FlashAttention และทวีตนี้ชี้ไปยังบทความวิจัยที่พวกเขาเผยแพร่ในวันนั้น: CODA

• หัวข้อบทความ: CODA: เขียนใหม่บล็อก Transformer เป็นโปรแกรม GEMM-Epilogue
• ที่อยู่ของเอกสารวิจัย: https://arxiv.org/abs/2605.19269
• ที่อยู่รหัส: https://github.com/HanGuo97/coda-kernels

ชื่อนี้ อ่านดูเหมือน “ฟินาเล่” แต่พูดออกมาเหมือน “CUDA” นักวิจัยจาก MIT, พรินซ์ตัน, Together AI และ Meta ได้พยายามใช้การดูดซับการเขียนโปรแกรมแบบใหม่ เพื่อจัดการกับการคำนวณเล็กๆ น้อยๆ ที่ถูกมองข้ามแต่ใช้เวลานานอย่างต่อเนื่องในกระบวนการฝึกสอน Transformer

ภาษีการขี้เกียจในการฝึกโมเดลขนาดใหญ่

เพื่อเข้าใจว่า CODA แก้ปัญหาอะไร ต้องเข้าใจก่อนว่าเวลาในการฝึกโมเดลขนาดใหญ่ใช้ไปกับอะไรบ้าง

เมื่อฝึกโมเดลขนาด 1 พันล้านพารามิเตอร์ในรูปแบบ LLaMA-3 บน NVIDIA H100 ผู้คนส่วนใหญ่มักจะมีความรู้สึกว่า: เวลาทั้งหมดใช้ไปกับการคูณเมทริกซ์และการคำนวณความสนใจ เพราะนั่นคือ “การคำนวณที่แท้จริง” ความรู้สึกนี้ถูกต้องในภาพรวม: การคูณเมทริกซ์ (GEMM) และความสนใจนั้นใช้พลังการประมวลผลหลัก

แต่ถ้าคุณเปิดตัววิเคราะห์ประสิทธิภาพแล้วดูอย่างละเอียด คุณจะพบว่ายังมี “โอเปอเรเตอร์ขนาดเล็ก” อีกชุดหนึ่งที่ค่อยๆ ใช้เวลาอยู่อย่างเงียบๆ: การทำให้เป็นมาตรฐาน (RMSNorm), ฟังก์ชันการกระตุ้น (SwiGLU, RoPE), การบวกแบบเหลือคงที่, การรวมข้ามชั้น... แม้แต่ละตัวจะมีปริมาณการคำนวณไม่มาก แต่กลับถูกย้ายเข้า-ออกหน่วยความจำกราฟิกบ่อยครั้งกับเทนเซอร์กลางขนาดใหญ่

นี่คือสิ่งที่เรียกว่า “ข้อจำกัดแบนด์วิธหน่วยความจำ”: เหมือนเชฟที่มีฝีมือยอดเยี่ยม แต่ต้องเดินไปเอาวัตถุดิบจากคลังสินค้าไกลๆ มาใช้ทีละอย่าง แล้วใช้เสร็จก็ต้องส่งคืนกลับไป แทนที่จะวางไว้บนโต๊ะทำงานใกล้มือ ไม่ว่ามือของเชฟจะเร็วแค่ไหน เวลาที่ต้องรอการขนส่งก็ยังเป็นการสูญเสียที่แท้จริง

แย่กว่านั้น ด้วยรูปแบบความแม่นยำต่ำของ NVIDIA เช่น FP8 และ FP4 ที่ทำให้การคำนวณเมทริกซ์เร็วขึ้น ต้นทุนสัมพัทธ์ของการดำเนินการ “ย้าย” เหล่านี้กลับเพิ่มขึ้น: การคูณเมทริกซ์เร็วขึ้น แต่ต้นทุนในการย้ายเทนเซอร์เข้าและออกไม่ได้ลดลงในสัดส่วนเดียวกัน

ในเอกสารวิจัยมีชุดข้อมูลหนึ่งที่ชัดเจน: เมื่อฝึกโมเดลพารามิเตอร์ 1 พันล้านบน H100 ด้วย TorchTitan การดำเนินการที่ไม่ใช่การคูณเมทริกซ์ใช้เวลาส่วนใหญ่ในระยะเวลาแบบ end-to-end และสัดส่วนนี้จะเด่นชัดยิ่งขึ้นเมื่อใช้ความแม่นยำ FP8

กรอบการเขียนโปรแกรมที่มีอยู่แทบไม่สามารถช่วยได้เลย PyTorch แสดงการคำนวณของ Transformer เป็นลำดับของโอเปอเรเตอร์ที่มีขอบเขตที่ชัดเจนระหว่างแต่ละตัว ขอบเขตเหล่านี้เหมาะสำหรับการหาอนุพันธ์อัตโนมัติ (autograd) แต่กลับขัดขวางการปรับรวมระหว่างโอเปอเรเตอร์: ขอบเขตของแต่ละโอเปอเรเตอร์มักจะเป็นการเขียนกลับหน่วยความจำที่ไม่จำเป็น

CODA: ซ่อนสมบัติไว้ใน「ตอนท้าย」

จุดเริ่มต้นของ CODA มาจากสังเกตที่เรียบง่าย

บน GPU นิวเคลียสการคูณเมทริกซ์ประสิทธิภาพสูง (GEMM) แบ่งออกเป็นสองส่วนหลัก: ลูปหลัก (mainloop) รับผิดชอบการคำนวณการคูณและบวกแบบแบ่งเมทริกซ์หลัก และส่วนท้าย (epilogue) รับผิดชอบการจัดการขั้นสุดท้ายก่อนบันทึกผลลัพธ์กลับไปยังหน่วยความจำของ GPU เช่น การบวกเบียส การแปลงชนิดข้อมูล และการปรับสเกลอย่างง่าย

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

การวิเคราะห์หลักของ CODA คือ: การดำเนินการที่ใช้หน่วยความจำมากใน Transformer หลายอย่างสามารถถูกปรับพารามิเตอร์ทางพีชคณิตและดำเนินการภายในหน้าต่าง "ท้ายสุด" นี้ได้

สิ่งนี้ต้องใช้ทักษะทางคณิตศาสตร์เล็กน้อย ตัวอย่างเช่นรูปแบบ GEMM-RMSNorm-GEMM ที่พบบ่อยที่สุด: ผลลัพธ์ของการคูณเมทริกซ์หนึ่งครั้ง จะผ่านการบวกแบบ residual, RMS normalization และจากนั้นจึงทำการคูณเมทริกซ์อีกครั้ง วิธีดั้งเดิมคือดำเนินการสามตัวดำเนินการแยกจากกันแบบเรียงลำดับ โดยผลลัพธ์กลางถูกบันทึกลงหน่วยความจำ GPU สองครั้ง

ทีม CODA พบว่า ปัจจัยการปรับขนาดแถว r ใน RMS normalization เนื่องจากเป็นสเกลาร์ที่ใช้ร่วมกันในแต่ละแถว จึงสามารถสลับลำดับได้กับการคูณเมทริกซ์ตามมา: สามารถเลื่อนการประยุกต์ใช้ r จาก “ก่อน GEMM ที่สอง” ไปเป็น “ช่วงท้ายของ GEMM ที่สอง” หลังจากเลื่อนแล้ว ช่วงท้ายของ GEMM ที่หนึ่งจำเป็นต้องคำนวณเฉพาะ “รากที่สองของค่าเฉลี่ยกำลังสองแบบย่อย” (partial RMS) เท่านั้น ซึ่งสามารถรวมเข้าด้วยกันโดยใช้เคอร์เนลลดขนาดที่มีน้ำหนักเบาอย่างยิ่ง ทำให้การคำนวณ RMSNorm แบบเต็มรูปแบบหายไป

การปรับพารามิเตอร์ใหม่ในลักษณะเดียวกันนี้ยังใช้ได้กับการดำเนินการอื่นๆ เช่น SwiGLU, RoPE (Rotation Position Encoding), และฟังก์ชันการสูญเสียแบบครอสเอนโทรปี รวมถึงยังใช้ได้กับการแพร่ย้อนกลับด้วย บทความวิจัยมีทฤษฎีบทพิสูจน์ว่า: หากส่วนท้ายของการคำนวณแบบก้าวหน้าเป็น “แบบแบ่งส่วนท้องถิ่น” การแพร่ย้อนกลับจะสืบทอดโครงสร้างเดียวกันโดยอัตโนมัติ โปรดเข้าไปดูรายละเอียดในบทความต้นฉบับ

ห้าประเภทของ「อิฐ」และชุด「ภาษาเลโก้」

CODA ไม่ใช่เคอร์เนลรวมที่เฉพาะเจาะจง แต่เป็นชุดของการดูดซับการเขียนโปรแกรม

มันยึดคงที่รอบหลักของ GEMM ที่ได้รับการปรับแต่งโดยผู้เชี่ยวชาญ แล้วเปิดเผยห้าประเภทของพื้นฐานที่สามารถรวมกันได้ในตำแหน่งท้าย:

• การแปลงแบบองค์ประกอบต่อองค์ประกอบ (การบวกแบบ residual, ฟังก์ชันการกระตุ้น, RoPE)
• การโหลดและเก็บเวกเตอร์ (การกระจายน้ำหนัก RMSNorm)
• การโหลดและจัดเก็บข้อมูลแบบแบ่งแมทริกซ์ (เก็บค่าการกระตุ้นชั่วคราวเพื่อใช้ในการย้อนกลับ)
• Block Reduction (Local RMS, Block Log-Sum-Exp)
• สถานะการเปลี่ยนแปลง (สถิติ max และ sum-exp ที่จำเป็นสำหรับการปรับมาตรฐานแบบออนไลน์)

ด้วยอิฐห้าประเภทนี้ สามารถครอบคลุมการดำเนินการส่วนใหญ่ทั้งหมดในกระบวนการการส่งผ่านข้างหน้าและข้างหลังของ Transformer มาตรฐาน ยกเว้นการให้ความสนใจ

ที่น่าสนใจยิ่งกว่านั้นคือความยืดหยุ่นของระบบนามธรรมนี้ต่อคำถามว่า “ใครควรเขียนโค้ด” ในการทดลอง บทความได้ประเมินรูปแบบการดำเนินการสองแบบ: แบบแรกคือโปรแกรมเมอร์มนุษย์เขียนขึ้นเอง และแบบที่สองคือใช้ Claude Code ในการสร้าง—โดยให้ AI สร้างโค้ดแกนหลักส่วนใหญ่จากคำอธิบายพื้นฐานของ CODA ตัวอย่างบางส่วน และบันทึกการดำเนินการ พร้อมการควบคุมดูแลแบบเบาๆ โดยมนุษย์

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

ผลการทดลอง

การทดสอบประสิทธิภาพของ CODA เลือกคู่ต่อสู้ที่ท้าทาย: cuBLAS ร่วมกับ torch.compile และ Liger Kernel และ FlashInfer ที่ได้รับการปรับแต่งโดยเฉพาะสำหรับ LLM

เอกสารวิจัยประเมินการใช้งานสองแบบสำหรับแต่ละเคอร์เนล: CODA (LLM) ซึ่งถูกสร้างโดย Claude Code โดยนักวิจัยให้คำอธิบายเกี่ยวกับพริมิทีฟ ตัวอย่างบางส่วน และบันทึกเทคนิคการดำเนินการที่อัปเดตอย่างต่อเนื่อง โดย AI สร้างโค้ดหลักและมนุษย์ทำการควบคุมแบบเบาๆ; CODA (Human) ถูกเขียนโดยโปรแกรมเมอร์มนุษย์อย่างอิสระ โดยใช้แนวคิดการรีพารามิเตอร์ระดับสูงเดียวกัน แต่ไม่พึ่งพาชุดพริมิทีฟของ CODA โดยตรง ผลลัพธ์จากทั้งสองกลุ่มถูกเปรียบเทียบกับไลบรารีที่ได้รับการปรับแต่ง เช่น cuBLAS + torch.compile, Liger Kernel และ FlashInfer

ในระดับโอเปอเรเตอร์เดียว โดยยกตัวอย่างรูปแบบทั่วไป GEMM-RMSNorm-GEMM CODA สามารถทำผลงาน vượtกว่าพื้นฐาน cuBLAS + PyTorch ภายใต้มิติซ่อนของโมเดลขนาด 1B, 7B และ 70B รูปแบบสุดท้ายอื่นๆ เช่น SwiGLU, RoPE และ cross-entropy ก็แสดงผลในลักษณะเดียวกัน

เคอร์เนลที่สร้างโดย LLM มีประสิทธิภาพใกล้เคียงกับเวอร์ชันที่เขียนด้วยมือโดยมนุษย์บนมาตรฐานส่วนใหญ่ และในบางการตั้งค่าอาจเหนือกว่าเล็กน้อย ซึ่งเป็นข้อสรุปที่พบได้ยากในสาขาการปรับแต่งเคอร์เนล GPU ที่มีอุปสรรคสูงมาโดยตลอด

ผลประโยชน์จากการย้อนกลับนั้นเด่นชัดเป็นพิเศษ: แกนย้อนกลับของ GEMM-Residual-PartialRMS-GEMM สามารถเร่งความเร็วได้ถึง 1.6 ถึง 1.8 เท่าเมื่อเทียบกับฐานอ้างอิง ส่วน SwiGLU ย้อนกลับก็มีการปรับปรุงประมาณ 1.4 ถึง 1.6 เท่า ในทิศทางนี้ ช่องว่างระหว่าง LLM กับการดำเนินการด้วยมือก็ยังเล็กเช่นกัน ซึ่งไม่น่าแปลกใจ: การย้อนกลับมีธรรมชาติที่เกี่ยวข้องกับการเข้าถึงเทนเซอร์กลางมากขึ้น จึงได้รับประโยชน์จากการรวมท้ายมากกว่า; และการออกแบบหน่วยพื้นฐานของ CODA มีความชัดเจนเพียงพอ ทำให้โมเดล AI สามารถจัดกลุ่มได้อย่างถูกต้อง

ในแบบจำลองแบบ end-to-end ของชั้น Transformer แบบเต็มรูปแบบ การเร่งความเร็วแบบ forward ของ CODA อยู่ที่ประมาณ 5% ถึง 20% สำหรับขนาดต่างๆ โดยมีประสิทธิภาพเด่นชัดยิ่งขึ้นในขนาดโมเดลที่ใหญ่ขึ้น (ซึ่งสอดคล้องกับมิติซ่อนที่ขนาด 70B)

ในด้านความแม่นยำของค่า ค่าพารามิเตอร์ซ้ำของ CODA ได้ปรับเวลาที่ใช้ปัจจัยการปรับขนาดของ RMSNorm แต่การทดลองแสดงให้เห็นว่าข้อผิดพลาดทางตัวเลขของมันอยู่ในระดับเดียวกับการใช้งานอ้างอิงของ PyTorch และในบางการตั้งค่า ข้อผิดพลาดยังน้อยกว่าอีก — เนื่องจากตัวสะสมที่มีความแม่นยำสูงกว่าในลูปหลักของ GEMM

CODA ทำได้อะไร: รายการเช็กเร็วเพื่อระบุขอบเขตความสามารถของ CODA ก่อนจะมองในมุมกว้างขึ้น

• ขอบเขต: การคำนวณเกือบทั้งหมดในการดำเนินการแบบฟอร์เวิร์ดและแบ็กเวิร์ดของ Transformer มาตรฐาน (เช่น สถาปัตยกรรม LLaMA) ยกเว้นการให้ความสนใจและการฝังคำ รวมถึง RMSNorm การบวกแบบรีซิดูอัล การกระตุ้น SwiGLU การเข้ารหัสตำแหน่งแบบ RoPE ความสูญเสียแบบครอสเอนโทรปี และการคำนวณเกรดเชิงย้อนกลับของการดำเนินการข้างต้น
• ผลเร่งความเร็ว: ในมิติซ่อนที่มีขนาดตั้งแต่ 1B ถึง 70B ระดับโอเปอเรเตอร์เดียวมีการปรับปรุงในระดับต่างๆ เมื่อเปรียบเทียบกับพื้นฐาน cuBLAS + torch.compile โดยเฉพาะอย่างยิ่งในกระบวนการย้อนกลับซึ่งได้รับผลประโยชน์อย่างชัดเจน (แกนบางตัวสามารถเร็วขึ้นกว่า 1.6 เท่าขึ้นไป); การเร่งความเร็วแบบ end-to-end สำหรับชั้น Transformer ทั้งหมดอยู่ที่ประมาณ 5% ถึง 20% โดยมีประสิทธิภาพเด่นชัดยิ่งขึ้นในโมเดลขนาดใหญ่
• ใครก็สามารถใช้: CODA ซึ่งดำเนินการโดยใช้ CuTeDSL (Python DSL ของ NVIDIA CUTLASS) รองรับวิธีการเขียนเคอร์เนลทั้งโดยโปรแกรมเมอร์มนุษย์และโมเดล AI โดยทั้งสองวิธีสามารถบรรลุประสิทธิภาพสูง
• ข้อจำกัดปัจจุบัน: ขณะนี้รองรับเฉพาะสภาพแวดล้อม GPU เดี่ยว ไม่เกี่ยวข้องกับการฝึกแบบกระจาย; การรีพารามิเตอร์มุ่งเน้นไปที่สถาปัตยกรรม Transformer มาตรฐาน ความเหมาะสมกับสถาปัตยกรรมอื่นๆ ยังต้องได้รับการยืนยัน

ข้อสรุป

CODA ไม่ใช่งานที่ทำคนเดียว มันเป็นการนำไปใช้จริงของแนวคิดประเภทหนึ่ง: บน GPU พื้นที่การปรับปรุงที่แท้จริงมักไม่อยู่ที่ “คำนวณอะไร” แต่อยู่ที่ “ย้ายอย่างไร”

FlashAttention ทำให้การคำนวณการให้ความสำคัญ “เข้าไปอยู่” ในหน่วยความจำบนชิป CODA พยายามทำให้การปกติและการใช้งานฟังก์ชันการกระตุ้นก็ “เข้าไปอยู่” เช่นกัน Triton ลดอุปสรรคในการเขียนเคอร์เนลที่กำหนดเอง ขณะที่ ThunderKittens, TileLang และอื่นๆ ได้สำรวจพื้นที่นี้ในระดับต่างๆ เพิ่มเติม งานเหล่านี้ร่วมกันชี้ไปในทิศทางเดียวกัน: การรวมความสะดวกในการแสดงกราฟโอเปอเรเตอร์ของ PyTorch เข้ากับประสิทธิภาพการดำเนินการที่ใกล้เคียงกับ CUDA ที่เขียนด้วยมือ ให้อยู่ในกรอบการทำงานที่สามารถเขียนโปรแกรมได้ชุดเดียว

ประโยคสุดท้ายของโพสต์ของ Tri Dao นั้น值得再回味: 「LLM และผู้เริ่มต้นสามารถเขียนโค้ร์ความเร็วแสงสำหรับการดำเนินการทั้งหมดของ Transformer」 ด้านหลังมีตรรกะที่ลึกซึ้งยิ่งขึ้น: เมื่อการออกแบบการดูดซับการเขียนโปรแกรมมีความเหมาะสมเพียงพอ โมเดล AI สามารถมีส่วนร่วมในการปรับปรุงโครงสร้างพื้นฐานการฝึกอบรมของตนเองได้ วงจรนี้เองคือจุดที่ CODA น่าสนใจที่สุด

ในมุมมองนี้ ชื่อ “CODA” อาจมีความหมายลึกซึ้งอีกประการหนึ่ง ในดนตรีคลาสสิก Coda คือส่วนที่สรุปบทเพลงทั้งหมดที่ปลายทาง ที่นี่ มันคือ “บทสรุป” ของแกน GEMM — และการเขียนบทสรุปนี้ให้ดีอาจเป็นบทสำคัญถัดไปในการเพิ่มประสิทธิภาพของระบบการฝึกอบรม Transformer