แนวคิดใหม่ในการเร่งความเร็วการฝึก Transformer: CODA เขียนการทำงานทั้งหมดใหม่เป็นการคูณเมทริกซ์ + บทสรุป หลุดพ้นจากคอขวดของแบนด์วิดท์หน่วยความจำ

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

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

• ชื่อเอกสาร: CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs
• ที่อยู่เอกสาร: https://arxiv.org/abs/2605.19269
• ที่อยู่โค้ด: https://github.com/HanGuo97/coda-kernels

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

พื้นหลัง: “ภาษีความขี้เกียจ” ของการฝึกโมเดลขนาดใหญ่

เพื่อทำความเข้าใจว่า CODA แก้ปัญหาอะไร ก่อนอื่นต้องรู้ว่าเวลาส่วนใหญ่ในการฝึกโมเดลขนาดใหญ่ไปอยู่ที่ไหน

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

อย่างไรก็ตาม หากคุณเปิด profiler และสังเกตอย่างละเอียด คุณจะพบว่ามี “operator ขนาดเล็ก” อีกกลุ่มหนึ่งที่ใช้เวลาอย่างเงียบๆ: Normalization (RMSNorm), ฟังก์ชันกระตุ้น (SwiGLU, RoPE), Residual Addition, Cross-layer Reduction… แต่ละตัวมีปริมาณการคำนวณไม่มาก แต่กลับย้าย tensor ขนาดใหญ่กลางทางเข้าและออกจาก VRAM บ่อยครั้ง

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

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

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

กรอบการเขียนโปรแกรมที่มีอยู่แทบจะทำอะไรไม่ได้กับเรื่องนี้ PyTorch แสดงการคำนวณของ Transformer เป็นลำดับของ operator ที่มีขอบเขตชัดเจนระหว่างกัน ขอบเขตนี้เป็นมิตรกับ autograd มาก แต่กลับขัดขวางการปรับให้เหมาะสมแบบ fusion ข้าม operator: ทุกขอบเขตของ operator มักจะเป็นการเขียนกลับไปยัง VRAM ที่ไม่จำเป็น

CODA: “Epilogue” ที่ซ่อนขุมทรัพย์

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

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

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

ข้อมูลเชิงลึกหลักของ CODA คือ: การดำเนินการที่ใช้หน่วยความจำเข้มข้นหลายอย่างใน Transformer จริงๆ แล้วสามารถถูก reparameterize ทางพีชคณิต และยัดเข้าไปในหน้าต่าง “epilogue” นี้เพื่อดำเนินการ

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

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

การ reparameterization ที่คล้ายกันนี้ใช้ได้กับ SwiGLU, RoPE (Rotary Position Embedding), Cross-Entropy Loss และแม้กระทั่งการแพร่กระจายย้อนกลับ (backpropagation) ในเอกสารมีทฤษฎีบทที่พิสูจน์ว่า: ถ้า forward epilogue เป็น “block-local” แล้ว backpropagation จะสืบทอดโครงสร้างเดียวกันโดยอัตโนมัติ รายละเอียดเฉพาะโปรดดูเอกสารต้นฉบับ

บล็อกห้าชนิดและ “ภาษาเลโก้” หนึ่งชุด

CODA ไม่ใช่เคอร์เนล fusion เฉพาะ แต่เป็นชุดนามธรรมการเขียนโปรแกรม

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

• การแปลงแบบ element-wise (residual addition, ฟังก์ชันกระตุ้น, RoPE)
• การโหลดและจัดเก็บเวกเตอร์ (broadcast น้ำหนัก RMSNorm)
• การโหลดและจัดเก็บเมทริกซ์แบบแบ่งส่วน (บันทึก activation กลางสำหรับ backpropagation)
• การ reduction แบบแบ่งส่วน (local root mean square, block log-sum-exp)
• การแปลงแบบมีสถานะ (สถิติ max และ sum-exp ที่จำเป็นสำหรับ online normalization)

ด้วยบล็อกห้าประเภทนี้ การดำเนินการเกือบทั้งหมด ยกเว้น Attention ในการ forward และ backward propagation ของ Transformer มาตรฐานสามารถครอบคลุมได้

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

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

ผลการทดลอง

เกณฑ์มาตรฐานของ CODA เลือกคู่แข่งที่ค่อนข้างโหด: cuBLAS บวกกับ torch.compile รวมถึง Liger Kernel และ FlashInfer ที่ปรับให้เหมาะสมสำหรับ LLM โดยเฉพาะ

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

ในระดับ operator เดียว ยกตัวอย่างรูปแบบทั่วไป GEMM-RMSNorm-GEMM CODA มีประสิทธิภาพเหนือกว่า baseline cuBLAS + PyTorch ในทุกมิติที่ซ่อนอยู่ซึ่งสอดคล้องกับขนาดโมเดล 1B, 7B และ 70B การผสมผสาน epilogue เช่น SwiGLU, RoPE, Cross-Entropy ก็มีประสิทธิภาพคล้ายกัน

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

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

เวอร์ชันเขียนใหม่เชิงลึกและลดการซ้ำซ้อน

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

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

รายการตรวจสอบความสามารถของ CODA

ก่อนที่จะพูดถึงมุมมองที่กว้างขึ้น เรามาชี้แจงขอบเขตการทำงานของ CODA ก่อน:

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

บทสรุป

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

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

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

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