Meta IKBO: วิธีลดความหน่วงในการอนุมานของระบบแนะนำลงหนึ่งลำดับความสำคัญด้วยการปรับปรุงการกระจายเสียงภายในเคอร์เนล?

มีความแตกต่างที่สำคัญแต่ถูกมองข้ามบ่อยครั้งระหว่างการอนุมานของระบบแนะนำและการอนุมานของ Large Language Model (LLM): เนื้อหาที่เป็นตัวเลือกจำนวนมากในคำขอเดียวกัน มักจะใช้ข้อมูลบริบทของผู้ใช้ชุดเดียวกัน ในการใช้งานแบบดั้งเดิม นั่นหมายความว่า “ฟีเจอร์ที่ใช้ร่วมกันในระดับคำขอ” เช่น เอมเบ็ดดิ้งของผู้ใช้หรือลำดับพฤติกรรมของผู้ใช้ จะต้องถูกคัดลอกซ้ำๆ เพื่อให้สอดคล้องกับขนาดแบตช์ของตัวเลือกและส่งเข้าไปในเลเยอร์ปฏิสัมพันธ์เพื่อคำนวณ การดำเนินการคัดลอกนี้ดูเหมือนเป็นเพียง “การกระจายเสียง” (broadcast) แต่ในการปรับใช้ระดับอุตสาหกรรม มันจะกินแบนด์วิดท์หน่วยความจำอย่างต่อเนื่อง ผลักดันต้นทุน I/O ให้สูงขึ้น และทำให้ความหน่วงแย่ลงเป็นเส้นตรงตามจำนวนตัวเลือกที่เพิ่มขึ้น

• In-Kernel Broadcast Optimization: Co-Designing Kernels for RecSys Inference
• https://pytorch.org/blog/in-kernel-broadcast-optimization-co-designing-kernels-for-recsys-inference/
• โค้ด: https://github.com/pytorch/FBGEMM/tree/main/fbgemm_gpu/experimental/ikbo
• ประมาณ 15,000 คำ ใช้เวลาอ่านประมาณ 56 นาที เวอร์ชันพอดแคสต์ประมาณ 17 นาที

Meta เสนอ In-Kernel Broadcast Optimization (IKBO) ในบทความนี้ การตัดสินใจหลักนั้นตรงไปตรงมา: การกระจายเสียงเป็นปัญหาการจัดเรียงข้อมูล (data layout problem) เป็นอันดับแรก ไม่ใช่ส่วนที่หลีกเลี่ยงไม่ได้ของการคำนวณ

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

สารบัญ

• คำถามสำคัญ
  • คำถามที่ 1: IKBO แทนที่คอขวดแบนด์วิดท์หน่วยความจำด้วยค่าใช้จ่ายในการกระจายการควบคุมและการซิงโครไนซ์หรือไม่?
  • คำถามที่ 2: การออกแบบแบบคู่ควบแน่นนี้แลกมาด้วยความสามารถในการประกอบ (composability) และกลายเป็นคอขวดในการวนซ้ำโมเดลหรือไม่?
• หนึ่ง การเพิ่มประสิทธิภาพการกระจายเสียงในเคอร์เนล: การกำจัดหน่วยความจำและการคำนวณที่ซ้ำซ้อน
  • 1.1 การกำจัดหน่วยความจำและการคำนวณที่ซ้ำซ้อน
  • 1.2 ประเภทของการเพิ่มประสิทธิภาพเคอร์เนล
  • 1.3 การออกแบบระบบแบบครบวงจร
  • 1.4 การเปรียบเทียบกับวิธีการอื่น
• สอง กรณีศึกษาเคอร์เนลเชิงลึกที่ 1: IKBO Linear Compression
  • 2.1 การแยกย่อยการคูณเมทริกซ์
  • 2.2 การเพิ่มประสิทธิภาพโครงร่างหน่วยความจำ
  • 2.3 การรวมการกระจายเสียงในเคอร์เนลสำหรับ Candidate GEMM
  • 2.4 การรวมหลายขั้นตอนแบบ Warp-Specialized โดยใช้ TLX
• สาม กรณีศึกษาเคอร์เนลเชิงลึกที่ 2: IKBO Flash Attention
  • 3.1 IKBO Flash Attention แก้ปัญหาคอขวด I/O ภายใต้เงื่อนไขขอบเขตของระบบแนะนำ
  • 3.2 การรวมเทคโนโลยีเคอร์เนลสมัยใหม่ (FA3, FA4) เข้ากับ IKBO บน TLX
  • 3.3 การรวม Self + Target Attention ผ่านการออกแบบร่วมของโมเดล
• สี่ การวัดประสิทธิภาพและสรุปผลลัพธ์
• ห้า บทสรุปและทิศทางในอนาคต
• เอกสารอ้างอิง
• ภาคผนวก
  • ภาคผนวก 1: การตั้งค่า Benchmark
  • ภาคผนวก 2: การวิเคราะห์ความเข้มข้นทางคณิตศาสตร์ (Arithmetic Intensity)
  • ภาคผนวก 3: การวิเคราะห์ผลลัพธ์โดยละเอียดของหัวข้อ 2.1
  • ภาคผนวก 4: วิธีการวิเคราะห์คอขวด
  • ภาคผนวก 5: การวิเคราะห์ผลลัพธ์โดยละเอียดของหัวข้อ 2.2
  • ภาคผนวก 6: การวิเคราะห์ผลลัพธ์โดยละเอียดของหัวข้อ 2.3
  • ภาคผนวก 7: โปรโตคอลการซิงโครไนซ์ Release-Acquire
  • ภาคผนวก 8: ตัวชี้วัด NCU Profiling ของ TLX และ Triton
  • ภาคผนวก 9: การวิเคราะห์ Roofline ของ Flash Attention ทั่วไปและ IKBO Flash Attention
  • ภาคผนวก 10: การใช้ SMEM ของ IKBO TLX FA3
  • ภาคผนวก 11: Benchmark ของ IKBO FA, CuTeDSL FA4 Hopper และ TLX FA3 Hopper ภายใต้เงื่อนไขขอบเขตของระบบแนะนำ
  • ภาคผนวก 12: Instruction Cache Miss ทำให้เกิดความหน่วงอย่างมีนัยสำคัญใน Consumer-2 Warpgroup

คำถามสำคัญ

คำถามที่ 1: IKBO แทนที่คอขวดแบนด์วิดท์หน่วยความจำด้วยค่าใช้จ่ายในการกระจายการควบคุมและการซิงโครไนซ์หรือไม่?

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

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

ประเด็นสำคัญคือการเปรียบเทียบขนาด ผู้เขียนระบุชัดเจนว่าภายใต้อัตราส่วนตัวเลือกต่อผู้ใช้จริงประมาณ 70:1 ความน่าจะเป็นที่จะเกิดพาธว่างนี้ “น้อยกว่า 0.7%” และค่าใช้จ่ายนั้นเล็กน้อย ในขณะที่ค่าใช้จ่ายในการคัดลอกหน่วยความจำที่ถูกกำจัดนั้นเป็นระบบและเพิ่มขึ้นเป็นเส้นตรงตามจำนวนตัวเลือก เช่น ในแผนเดิม ต้องคัดลอกฟีเจอร์ผู้ใช้ประมาณ 70 ครั้ง ใช้ปริมาณการรับส่งหน่วยความจำระดับกลางถึง 0.87 GB

ปรัชญาการออกแบบที่ลึกซึ้งยิ่งขึ้นคือ: ความเชี่ยวชาญของ warp ใน TLX ไม่ได้หลีกเลี่ยงความแตกต่างของการควบคุม แต่ทำให้มันชัดเจนและควบคุมได้ โดยการแบ่ง CTA ออกเป็นกลุ่มผู้ผลิตและผู้บริโภค และจัดการการซิงโครไนซ์อย่างแม่นยำด้วย named barrier IKBO จะเปลี่ยนการแตกกิ่งที่คาดเดาไม่ได้ให้เป็นไปป์ไลน์ที่มีโครงสร้าง ข้อมูล NCU แสดงให้เห็นว่าแม้จะมีการซิงโครไนซ์ ปริมาณงาน DRAM ก็เพิ่มขึ้นจาก 39% เป็น 52% และการใช้ประโยชน์คอขวด L2 เพิ่มขึ้นเป็น 84%

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

คำถามที่ 2: การออกแบบแบบคู่ควบแน่นนี้แลกมาด้วยความสามารถในการประกอบ (composability) และกลายเป็นคอขวดในการวนซ้ำโมเดลหรือไม่?

IKBO อ้างว่าบรรลุการเร่งความเร็วอย่างมีนัยสำคัญผ่าน “การออกแบบร่วมของเคอร์เนล-โมเดล-ระบบ” และทำให้การออกแบบโมเดลเป็นอิสระจากข้อจำกัด แต่ความยืดหยุ่นนี้แลกมาด้วยหลักการวิศวกรรมซอฟต์แวร์ของ “localality” และ “composability” หรือไม่? เมื่อจำเป็นต้องดีบักปัญหาความแม่นยำหรือเปลี่ยนโครงสร้างโมเดลอย่างรวดเร็ว การเพิ่มประสิทธิภาพแบบคู่ควบแน่นนี้จะกลายเป็นคอขวดของความเร็วในการวนซ้ำหรือไม่?

ผู้เขียนแสดงให้เห็นถึงแนวทางการออกแบบที่ตรงกันข้าม: การปกป้องโมดูลาร์ผ่านสถาปัตยกรรมระบบ ไม่ใช่ผ่านธรรมเนียมปฏิบัติ

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

เกี่ยวกับปัญหาความแม่นยำเชิงตัวเลข ภาคผนวก 3 ตอบสนองโดยตรงต่อความแตกต่างเชิงตัวเลขที่อาจเกิดขึ้นจากการแยกย่อย GEMM โดยระบุอย่างชัดเจนว่า “เป็นไปตามมาตรฐานความแม่นยำ” ซึ่งแสดงให้เห็นว่าทีมงานตระหนักถึงปัญหาที่อาจเกิดขึ้นนี้และเชื่อว่าความเสี่ยงอยู่ในขอบเขตที่ควบคุมได้ สำหรับข้อกังวลว่าความหมายของโอเปอเรเตอร์จะถูกทำลาย IKBO ไม่ได้แนะนำโอเปอเรเตอร์ใหม่ แต่เปลี่ยนกลไกการใช้งานภายในของโอเปอเรเตอร์ทั่วไปที่มีอยู่ เมื่อนักวิจัยต้องการเปลี่ยนฟังก์ชันกระตุ้น พวกเขาไม่จำเป็นต้องแก้ไขโค้ดใดๆ ทีมเคอร์เนลเพียงแค่ต้องบำรุงรักษาไลบรารีตัวแปร IKBO อย่างอิสระ

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

หนึ่ง การเพิ่มประสิทธิภาพการกระจายเสียงในเคอร์เนล: การกำจัดหน่วยความจำและการคำนวณที่ซ้ำซ้อน

บทความนี้แนะนำ In-Kernel Broadcast Optimization (IKBO): วิธีการออกแบบร่วมของเคอร์เนล-โมเดล-ระบบสำหรับการอนุมานโมเดลแนะนำ เพื่อกำจัดการกระจายเสียงเอมเบ็ดดิ้งผู้ใช้ที่ซ้ำซ้อน ในระบบแนะนำระดับการผลิต สำหรับคำขอที่กำหนด เอมเบ็ดดิ้งผู้ใช้จะเหมือนกันทุกประการในรายการตัวเลือกทั้งหมด แต่แนวทางปฏิบัติมาตรฐานยังคงต้องการการคัดลอกอย่างชัดเจน ทำให้เกิดการสิ้นเปลืองแบนด์วิดท์หน่วยความจำและการคำนวณที่เพิ่มขึ้นเป็นเส้นตรงตามจำนวนตัวเลือกที่เพิ่มขึ้น ข้อมูลเชิงลึกหลักของ IKBO คือ: การกระจายเสียงเป็นปัญหาการจัดเรียงข้อมูลโดยพื้นฐาน ไม่ใช่สิ่งที่จำเป็นสำหรับการคำนวณ เคอร์เนล IKBO แต่ละตัวรับอินพุตขนาดแบตช์ผู้ใช้และขนาดแบตช์ตัวเลือกที่ไม่ตรงกันโดยธรรมชาติ และจัดการการกระจายเสียงภายในเคอร์เนล ดังนั้นจึงไม่สร้างเทนเซอร์ที่ถูกคัดลอกขึ้นมาจริง บทความนี้แสดงให้เห็นวิธีการนี้ผ่านสองกรณีศึกษาเคอร์เนลเชิงลึก: Linear Compression และ Flash Attention

IKBO ถูกปรับใช้ในสแต็กการอนุมานระบบแนะนำของ Meta ครอบคลุมตั้งแต่โมเดลการจัดอันดับ早期ไปจนถึง晚期 รองรับ GPU และ MTIA (Meta Training and Inference Accelerator) บนโมเดลที่ออกแบบร่วมกัน มันสามารถลดความหน่วงของเครือข่ายที่เน้นการคำนวณได้มากถึงสองในสาม นอกจากนี้ยังเป็นกระดูกสันหลังของ Meta Adaptive Ranking Model [1] ซึ่งเป็นเฟรมเวิร์กที่เน้นคำขอและมีประสิทธิภาพในการอนุมาน สำหรับการให้บริการโมเดลขนาด LLM ในสภาพแวดล้อมการผลิต

บน H100 SXM5 เคอร์เนล Linear Compression IKBO ของเราผ่านสี่ขั้นตอนการออกแบบร่วมแบบก้าวหน้า: การแยกย่อยการคูณเมทริกซ์ การจัดแนวหน่วยความจำ การรวมการกระจายเสียง และการรวมหลายขั้นตอนแบบ warp-specialized ผ่าน TLX (Triton Low-Level Extensions) [2] ทำให้ได้ความเร่งประมาณ 4 เท่า สำหรับ Flash Attention IKBO ให้ปริมาณงานเพิ่มขึ้น 2.4×/6.4× (เฉพาะเคอร์เนล Attention/เคอร์เนล Attention บวกค่าใช้จ่ายในการกระจายเสียง) เมื่อเทียบกับ CuTeDSL FA4-Hopper ที่ไม่ได้ออกแบบร่วมกัน โดยทำได้ถึง 621 BF16 TFLOPs

แตกต่างจากวิธีการ “เลี่ยงการคัดลอก” เช่น การกระจายเสียงระดับระบบหรือการแยกเครือข่าย IKBO กำจัดการคัดลอกโดยตรงในระดับดั้งเดิมของการคำนวณ โดยให้คุณภาพปฏิสัมพันธ์ที่หนาแน่นด้วยต้นทุนที่เกือบจะเป็นอิสระ

หมายเหตุ: งานนี้เสร็จสมบูรณ์ในขณะที่ผู้เขียนทำงานที่ Meta

1.1 การกำจัดหน่วยความจำและการคำนวณที่ซ้ำซ้อน

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

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

ฟีเจอร์ทั้งสองประเภทจะผ่านการค้นหาเอมเบ็ดดิ้ง (embedding lookup) และการประมวลผลตามมา เพื่อสร้างการแสดงผลแบบเอมเบ็ดดิ้ง ในตำแหน่งต่างๆ ของโมเดล เลเยอร์ปฏิสัมพันธ์ (เช่น การฉายภาพเชิงเส้น การข้ามฟีเจอร์ Target Attention) จะรวมเอมเบ็ดดิ้งผู้ใช้กับเอมเบ็ดดิ้งตัวเลือกเข้าด้วยกัน เราจะเรียกเอมเบ็ดดิ้งที่ใช้ร่วมกันในทุกตัวเลือกของคำขอว่า Request-Only (RO) และเรียกเอมเบ็ดดิ้งเฉพาะแต่ละตัวเลือกว่า Non-Request-Only (NRO)

รูปที่ 1. กระแสข้อมูลการอนุมานระบบแนะนำแบบง่ายอย่างยิ่ง เอมเบ็ดดิ้งผู้ใช้แบบ Request-Only (RO) จะต้องถูกกระจายเสียง (คัดลอก) เพื่อให้ตรงกับมิติแบตช์ของรายการตัวเลือกแบบ Non-Request-Only (NRO) ก่อนที่จะเข้าสู่เลเยอร์ปฏิสัมพันธ์ IKBO กำจัดการสร้างขึ้นมาจริงนี้โดยจัดการการกระจายเสียงภายในแต่ละเคอร์เนล รูปนี้เปรียบเทียบวิธีการดั้งเดิมกับแผน IKBO อย่างชัดเจน ในวิธีการดั้งเดิม การกระจายเสียงเป็นขั้นตอนการเตรียมข้อมูลภายนอกที่ชัดเจน ในขณะที่ IKBO ทำให้มันเป็นส่วนหนึ่งของกระบวนการคำนวณ นี่ไม่ใช่แค่การเปลี่ยนตำแหน่งการดำเนินการ แต่หมายถึงโครงสร้างกราฟการคำนวณที่เรียบง่ายขึ้น และรันไทม์ไม่จำเป็นต้องจัดสรรและถ่ายโอนหน่วยความจำสำหรับเทนเซอร์ที่ถูกคัดลอกขนาดใหญ่อีกต่อไป

เลเยอร์ปฏิสัมพันธ์ต้องการให้เทนเซอร์มีมิติแบตช์ที่ตรงกัน ในแบตช์ที่ให้บริการผู้ใช้ประมาณ 15 คน แต่มีตัวเลือก 1024 รายการ เอมเบ็ดดิ้ง RO จะต้องถูกกระจายเสียงก่อน นั่นคือคัดลอกประมาณ 70 ครั้ง เพื่อให้สอดคล้องกับขนาดแบตช์ NRO ก่อนที่ปฏิสัมพันธ์ใดๆ จะเกิดขึ้น (รูปที่ 1) ในขณะที่สถาปัตยกรรมโมเดลพัฒนาจาก DLRM [1] และ DCN [2] ไปเป็นโมเดลลำดับเช่น HSTU [3] และ Phoenix [4] ของ X ปฏิสัมพันธ์ระหว่างผู้ใช้และตัวเลือกก็เพิ่มขึ้นอย่างต่อเนื่อง แต่ปฏิสัมพันธ์ที่สมบูรณ์ยิ่งขึ้นก็หมายถึงต้นทุนที่สูงขึ้นเช่นกัน: ฟีเจอร์ผู้ใช้ต้องถูกกระจายเสียงไปยังตัวเลือกทั้งหมด ในการอนุมาน เมื่อขนาดแบตช์อยู่ในช่วง 10 ถึง 10,000+ ค่าใช้จ่ายในการคัดลอกนี้ทำให้เกิดต้นทุนการคำนวณและหน่วยความจำที่สำคัญ และเพิ่มขึ้นเป็นเส้นตรงตามจำนวนตัวเลือก

การกระจายเสียงเป็นปัญหาการจัดเรียงข้อมูล ไม่ใช่สิ่งที่จำเป็นสำหรับการคำนวณ การมองโมเดลและระบบอนุมานใหม่ด้วยมุมมองนี้ จะเผยให้เห็นพื้นที่สำหรับการเพิ่มประสิทธิภาพในทุกเลเยอร์: รันไทม์การอนุมานกำจัดการกระจายเสียงระดับระบบ เลเยอร์โมเดลที่เกี่ยวข้องกับผู้ใช้เท่านั้นทำงานบนขนาดแบตช์ผู้ใช้ที่เล็กกว่า ในขณะที่เคอร์เนลที่จัดการข้อมูลทั้งสองประเภทได้รับการออกแบบใหม่เพื่อจัดการการกระจายเสียงภายใน ดังนั้นจึงไม่มีเทนเซอร์ที่ถูกคัดลอกใดๆ ถูกสร้างขึ้นมาจริง IKBO ถูกปรับใช้ในสแต็กการอนุมานระบบแนะนำของ Meta ตั้งแต่การจัดอันดับ早期ไปจนถึง晚期 ครอบคลุม GPU และ MTIA บนโมเดลที่ออกแบบร่วมกัน สามารถลดความหน่วงของเครือข่ายที่เน้นการคำนวณได้มากถึงสองในสาม

บทความนี้มุ่งเน้นไปที่เลเยอร์เคอร์เนล โดยขยายความผ่านสองกรณีศึกษาเชิงลึก: Linear Compression และ Flash Attention

1.2 ประเภทของการเพิ่มประสิทธิภาพเคอร์เนล

• Type I — การดำเนินการที่แยกย่อยได้ โดยการปรับโครงสร้างทางคณิตศาสตร์ ส่วน Request-Only (RO) สามารถคำนวณอย่างอิสระภายใต้ขนาดแบตช์เล็ก และรวมกับส่วน Non-Request-Only (NRO) ในตอนท้ายเท่านั้น ซึ่งช่วยประหยัดทั้งแบนด์วิดท์หน่วยความจำและปริมาณการคำนวณ
• Type II — การเพิ่มประสิทธิภาพหน่วยความจำล้วนๆ จัดการการกระจายเสียง RO-NRO ภายในเคอร์เนล หลีกเลี่ยงการเคลื่อนย้ายข้อมูลที่ซ้ำซ้อน ทำให้เคอร์เนลหลุดพ้นจากสถานะที่ถูกจำกัดด้วย I/O

1.3 การออกแบบระบบแบบครบวงจร

การปรับใช้ IKBO ส่งผลกระทบต่อสามเลเยอร์ของสแต็กโครงสร้างพื้นฐาน:

1. เคอร์เนล: เคอร์เนล GPU แบบกำหนดเองที่ยอมรับขนาดแบตช์ RO/NRO ที่ไม่ตรงกัน และจัดการการกระจายเสียงภายใน (ดูหัวข้อ 2 และ 3 ด้านล่าง)
2. ข้อกำหนดการคอมไพล์: คอมไพเลอร์ ML ต้องการช่วงรูปร่างไดนามิกสำหรับโอเปอเรเตอร์แต่ละตัว เพื่อเลือกเคอร์เนลที่ตรงกับรูปร่าง เมื่อมีขนาดแบตช์เดียว ก็ง่าย เมื่อมีสองขนาด (ผู้ใช้และตัวเลือก) หรือมากกว่า การตัดสินใจอย่างน่าเชื่อถือว่าโอเปอเรเตอร์แต่ละตัวใช้มิติแบตช์ใด—โดยเฉพาะอย่างยิ่งเมื่อปฏิสัมพันธ์ทำให้แหล่งที่มาของแบตช์ไม่ชัดเจน—จำเป็นต้องมีระบบอัตโนมัติที่เป็นระบบ
3. การอนุมาน: รันไทม์ส่งผ่านการแมป candidate-to-user ไปยังโมเดล แทนที่จะสร้างผลลัพธ์การกระจายเสียงขึ้นมาก่อน

เคอร์เนลเหล่านี้เข้าสู่โมเดลได้สองวิธี:

1. การนำไปใช้โดยตรง: ผู้เขียนโมเดลรวมเคอร์เนล IKBO เข้ากับนิยามโมเดลโดยตรง เมื่ออัตราส่วน candidate-to-user ระหว่างการฝึกมากกว่า 1 เคอร์เนลเดียวกันก็สามารถลดต้นทุนการฝึกได้เช่นกัน
2. การแปลงระหว่างการอนุมาน: พาส (pass) จะแทนที่โอเปอเรเตอร์มาตรฐานด้วยโอเปอเรเตอร์ IKBO ที่เทียบเท่าโดยอัตโนมัติระหว่างการอนุมาน โดยไม่ต้องแก้ไขโค้ดโมเดล

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

1.4 การเปรียบเทียบกับวิธีการอื่น

วิธีการที่มีอยู่มักจะหลีกเลี่ยงการกระจายเสียง มากกว่าที่จะกำจัดมันอย่างแท้จริง

1. การกระจายเสียงระดับระบบจะสร้างเทนเซอร์ที่ถูกคัดลอกขึ้นมาก่อนที่จะส่ง GPU dispatch — ใช้งานง่าย แต่มีต้นทุนสูง และต้นทุนเพิ่มขึ้นเป็นเส้นตรงตามจำนวนตัวเลือก
2. การแยกเครือข่าย (ROO) [5] แบ่งโมเดลออกเป็นเครือข่ายย่อย RO และ NRO ซึ่งสามารถลดงานที่ซ้ำซ้อนได้ แต่จะจำกัดตำแหน่งที่ปฏิสัมพันธ์ระหว่างผู้ใช้และตัวเลือกสามารถเกิดขึ้นได้ และยังคงทำให้เกิดต้นทุนเพิ่มเติมภายใต้ขนาดแบตช์ RO ที่เล็กกว่า

ทั้งสองวิธียังคงรักษาการกระจายเสียงในรูปแบบของ “เทนเซอร์ที่ถูกสร้างขึ้นมา” IKBO กำจัดมันโดยตรงในระดับดั้งเดิมของการคำนวณ: การประหยัดจะขยายตัวตามอัตราส่วน candidate-to-user ที่เพิ่มขึ้น รูปแบบปฏิสัมพันธ์ใดๆ ก็ไม่ต้องแบกรับต้นทุนการกระจายเสียง และมิติแบตช์ NRO ที่สมบูรณ์ในเคอร์เนลแบบรวมยังให้ occupancy ของ GPU

IKBO ถูกปรับใช้บนตัวเร่ง GPU และ MTIA บทความนี้มุ่งเน้นไปที่การออกแบบเคอร์เนล H100 GPU เพื่ออธิบายหลักการเพิ่มประสิทธิภาพหลัก

สอง กรณีศึกษาเคอร์เนลเชิงลึกที่ 1: IKBO Linear Compression

2.1 การแยกย่อยการคูณเมทริกซ์

LCE พื้นฐานจะดำเนินการคูณเมทริกซ์แบบแบตช์ (batched matmul) หนึ่งครั้งสำหรับตัวเลือก B ทั้งหมด เอมเบ็ดดิ้งอินพุตจะต่อส่วนผู้ใช้และส่วนตัวเลือกเข้าด้วยกันในมิติ K — อย่างไรก็ตาม สำหรับตัวเลือกทั้งหมดของผู้ใช้คนเดียวกัน เอมเบ็ดดิ้งผู้ใช้นั้นเหมือนกันทุกประการ

รูปที่ 2 แสดงแนวคิดหลักของการแยกย่อย LCE: มุมซ้ายบนคือการคูณเมทริกซ์แบบแบตช์พื้นฐาน มุมขวาบนคือการแยกเอมเบ็ดดิ้งตามมิติ K และการทำ unique ผู้ใช้ ด้านล่างแสดงการดำเนินการ GEMM สองแบบอิสระ ซึ่งในที่สุดจะรวมเอาต์พุตที่ถูกบีบอัดผ่านการบวกแบบกระจายเสียง รูปนี้แสดงให้เห็นถึงการเปลี่ยนแปลงกระบวนทัศน์จาก “คัดลอกก่อนแล้วค่อยคำนวณ” เป็น “คำนวณก่อนแล้วค่อยรวมแบบเบา” ขั้นแรกโดยการแยกเอมเบ็ดดิ้งของผู้ใช้และสินค้าตัวเลือก ขนาดของการคูณเมทริกซ์ฝั่งผู้ใช้จะลดลงหลายสิบเท่าทันที ขั้นตอนที่สองเลื่อนการดำเนินการกระจายเสียงจากเวกเตอร์เอมเบ็ดดิ้งมิติสูงเต็มรูปแบบไปยังผลลัพธ์ที่ถูกบีบอัดมิติต่ำ ซึ่งช่วยลดปริมาณข้อมูลที่ต้องจัดการในการกระจายเสียงลงอย่างมาก

เลื่อนการกระจายเสียงไปหลังจากการคูณเมทริกซ์ เนื่องจากน้ำหนัก W ไม่ขึ้นกับแบตช์ เราจึงใช้ประโยชน์จากคุณสมบัติเชิงเส้นในการแยกย่อย: แยกบล็อกเอมเบ็ดดิ้งผู้ใช้และบล็อกเอมเบ็ดดิ้งตัวเลือกตามมิติ K ทำ unique เอมเบ็ดดิ้งผู้ใช้ที่ซ้ำกัน จากนั้นดำเนินการ GEMM สองแบบอิสระบนขนาดแบตช์ตามธรรมชาติของแต่ละส่วน ด้วยวิธีนี้ เอมเบ็ดดิ้งผู้ใช้ไม่จำเป็นต้องถูกคัดลอกก่อน matmul อีกต่อไป แค่กระจายเสียงผลลัพธ์ที่ถูกบีบอัดซึ่งมีขนาดเล็กกว่าเท่านั้น ดูรายละเอียดในรูปที่ 2 ในสถานการณ์ทั่วไปที่อัตราส่วน candidate-to-user อยู่ที่ประมาณ 70 ขนาดแบตช์ผู้ใช้จะลดลงอย่างรวดเร็วจาก B=1024 เหลือ B_user ≈ 15 ทำให้ปริมาณการคำนวณฝั่งผู้ใช้ลดลง 70 เท่า กระบวนการแยกย่อยนี้ใช้ PyTorch มาตรฐานทั้งหมด

ผลลัพธ์: ความหน่วงลดลงจาก 1.944 ms เป็น 1.389 ms (ลดลง 28.5% รายละเอียดการตั้งค่า benchmark ดูในภาคผนวก 1) ความเข้มข้นทางคณิตศาสตร์ของ GEMM แบบแบตช์ดั้งเดิมอยู่ที่ประมาณ 356 FLOPs/Byte ซึ่งต่ำกว่าจุดสมดุลของเครื่อง H100 ที่ประมาณ 495 FLOPs/Byte (รายละเอียดการ推导见ภาคผนวก 2) GEMM ทั้งสองแบบหลังการแยกย่อยยังคงถูกจำกัดด้วยหน่วยความจำ ดังนั้นการปรับปรุงประสิทธิภาพหลักมาจากการลดต้นทุนหน่วยความจำ การทำ unique ช่วยลดต้นทุนหน่วยความจำลงมากกว่าครึ่งหนึ่ง—เนื่องจาก GEMM ฝั่งผู้ใช้ (B_user ≈ 15 แทนที่จะเป็น B = 1024) แทบจะไม่มีนัยสำคัญ

ควรชี้ให้เห็นเป็นพิเศษว่าการแยกย่อยที่นี่เพียงแค่ย้ายการดำเนินการกระจายเสียงไปหลัง matmul: แทนที่จะคัดลอกเอมเบ็ดดิ้งอินพุตมิติ K เต็มรูปแบบก่อน GEMM ตอนนี้แค่กระจายเสียงผลลัพธ์ที่ถูกบีบอัดซึ่งมีขนาดเล็กกว่า ดังนั้นต้นทุนจึงต่ำกว่ามาก หัวข้อ 2.3 จะกำจัดการกระจายเสียงที่เหลืออยู่นี้ให้หมดไปโดยการรวมการกระจายเสียงในเคอร์เนล

คอขวดปัจจุบัน主要集中在 L1/TEX pipeline utilization (84%) มากกว่า DRAM utilization — ความไม่สมดุลนี้值得警惕 เราจะวิเคราะห์โดยละเอียดในหัวข้อถัดไป ผลการ profiling ที่ละเอียดยิ่งขึ้นสามารถดูได้ในภาคผนวก 3

2.2 การเพิ่มประสิทธิภาพโครงร่างหน่วยความจำ

การวิเคราะห์ผลลัพธ์โดยละเอียดของ GEMM หลังการแยกย่อยเผยให้เห็นสภาวะที่ไม่สมดุล: L1/TEX ถึง 84% ของ峰值 ในขณะที่ DRAM อยู่ที่เพียง 19% ซึ่งบ่งชี้ว่าความกว้างของการโหลดหน่วยความจำแคบเกินไปโดยไม่จำเป็น โค้ด SASS ยืนยัน: แต่ละคำสั่ง cp.async คัดลอกเพียง 4 ไบต์ แทนที่จะโหลดแบบ 128-bit

LDGSTS.E.LTC128B P0, [R203], [R38.64] // 4 bytes
LDGSTS.E.LTC128B P1, [R203+0x4], [R38.64+0x4] // 4 bytes (รวม ×4 ก็แค่ 16B)

ความกว้างของ cp.async ถูกจำกัดโดยข้อจำกัดการจัดแนวตามธรรมชาติของพอยน์เตอร์ต้นทาง เมทริกซ์ A ใช้โครงร่าง row-major (M, K) โดยมี stride เป็น K × 2 ไบต์ ดังนั้นเมื่อ K ไม่ใช่倍数ของ 8 stride นี้จะทำลายการจัดแนว 128-bit

ข้อมูลเชิงลึกของการออกแบบร่วมโมเดล-เคอร์เนล: การจัดแนวหน่วยความจำ本身เป็นปัญหาการเพิ่มประสิทธิภาพ GPU แบบคลาสสิก แต่การดำเนินการแยกย่อยทำให้มันกลายเป็นความท้าทายของการออกแบบร่วมโมเดล-เคอร์เนล K มาจาก torch.cat ของเทนเซอร์เอมเบ็ดดิ้งหลายตัว และขนาดของเทนเซอร์เหล่านี้ขึ้นอยู่กับปัจจัยการกำหนดค่าโมเดลหลายประการ หลังจากแยกย่อยแล้ว แทบจะเป็นไปไม่ได้เลยที่จะปรับปัจจัยเหล่านี้ด้วยตนเองเพื่อให้แน่ใจว่าเอมเบ็ดดิ้งที่แยกย่อยแล้วยังคงจัดแนว倍数ได้อย่างสมบูรณ์แบบ ดังนั้นจึงจำเป็นต้องมีโซลูชันที่เป็นระบบ

โซลูชัน: ขยาย K ที่แยกย่อยแต่ละตัวโดยการเติมศูนย์ (pad) ไปยัง倍数ถัดไปของ 8 เราพิสูจน์ว่าในการส่ง forward และ backward การดำเนินการนี้เทียบเท่าทางคณิตศาสตร์กับการคำนวณดั้งเดิม (ดูข้อพิสูจน์ 1 ด้านล่าง) และด้วยความช่วยเหลือของตัววางแผนหน่วยความจำของ ML คอมไพเลอร์ นี่เป็นเพียงการคัดลอกค่าคงที่ต้นทุนต่ำเท่านั้น

สูตรอธิบายเทคนิคการ Padding ในการคูณเมทริกซ์ โดยมีเป้าหมายเพื่อปรับมิติเมทริกซ์ให้เป็น倍数ที่เหมาะสมกับฮาร์ดแวร์ (ในที่นี้คือ倍数ของ 8) เพื่อเพิ่มประสิทธิภาพการคำนวณ ในขณะที่รับประกันความถูกต้องทางคณิตศาสตร์ทั้งในการส่ง forward และ backward ในการตั้งค่า การคูณเมทริกซ์ โดยที่ , , . เมื่อ Padding เติม คอลัมน์ศูนย์ให้ , เติม แถวศูนย์ให้แต่ละแบตช์ของ ทำให้มิติใหม่ เป็น倍数ของ 8 ในการส่ง forward ผลลัพธ์ของ เหมือนกับการคูณเมทริกซ์ดั้งเดิม ทุกประการ การเติมศูนย์ไม่มีผลต่อผลลัพธ์การคำนวณ ในการส่ง backward เกรเดียนต์ของ loss ต่อ จะมีผลเฉพาะ คอลัมน์แรกเท่านั้น ส่วนคอลัมน์ศูนย์ที่เหลือมีเกรเดียนต์เป็น 0 เกรเดียนต์ของ loss ต่อ จะมีผลเฉพาะ แถวแรกเท่านั้น ส่วนแถวศูนย์ที่เหลือมีเกรเดียนต์เป็น 0 ทั้งสองสามารถ还原เกรเดียนต์ของเมทริกซ์ดั้งเดิมได้อย่างถูกต้อง วิธีการ Padding นี้在不

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

☕ สนับสนุนค่ากาแฟทีมงาน

หากคุณชอบบทความนี้ สามารถสนับสนุนเราได้ผ่าน PromptPay

SCAN TO PAY WITH ANY BANK