การเปิดใช้งานโมเดลระดับล้านล้านพารามิเตอร์บน AWS EFA

ที่ Perplexity เราใช้โมเดลที่ดีที่สุดสำหรับผลิตภัณฑ์ API และทีมวิจัยของเรา โมเดล Mixture-of-Experts ขนาดใหญ่แบบโอเพนซอร์ส เช่น Kimi-K2 ก่อให้เกิดความท้าทายเฉพาะตัว เนื่องจากโหนดการอนุมานขนาดใหญ่ที่สุดที่มี GPU NVIDIA H200 จำนวน 8x ไม่สามารถรองรับโมเดลเหล่านี้ได้อย่างมีประสิทธิภาพ จึงจำเป็นต้องมีการปรับใช้งานแบบหลายโหนด เรานำเสนอชุดเคอร์เนลสำหรับขนานผู้เชี่ยวชาญ (expert parallelism) ที่บรรลุเวลาแฝงระดับล้ำสมัยบน ConnectX-7 ซึ่งเหนือกว่าประสิทธิภาพของ DeepEP นอกจากนี้ เคอร์เนลชุดเดียวกันนี้ยังเป็นชุดแรกที่บรรลุเวลาแฝงที่ใช้งานได้จริงบน AWS Elastic Fabric Adapter (EFA) ซึ่งช่วยให้สามารถปรับใช้งานโมเดลที่มีพารามิเตอร์ระดับล้านล้านได้

ลองใช้เคอร์เนลของเราบน GitHub และอ่านงานวิจัยฉบับเต็มได้ที่ arXiv

บทนำ

สถาปัตยกรรม Mixture-of-Experts (MoE) กลายเป็นสถาปัตยกรรมมาตรฐานสำหรับการปรับขนาดไปสู่พารามิเตอร์ระดับหลายแสนล้านหรือล้านล้านตัว พร้อมทั้งยังให้เวลาแฝงที่เหมาะสมในขณะทำการอนุมาน โมเดล MoE จะแทนที่เลเยอร์หนาแน่นของทรานสฟอร์เมอร์ด้วยชุดผู้เชี่ยวชาญและเลเยอร์การกำหนดเส้นทางเชิงเส้นขนาดเล็ก ซึ่งทำหน้าที่ตัดสินใจว่าควรนำโทเค็นไปคูณกับผู้เชี่ยวชาญใด การจัดกลุ่มโทเค็นที่เกี่ยวข้องกับผู้เชี่ยวชาญเดียวกันจะช่วยให้ปัญหา กลายเป็นแบบเบาบาง (sparse) และสามารถขนานได้อย่างง่ายดาย เนื่องจากตอนนี้ผู้เชี่ยวชาญที่แตกต่างกันสามารถโฮสต์อยู่บน GPU ที่แตกต่างกันซึ่งกระจายอยู่ทั่วหลายโหนดได้

ความแตกต่างจากรูปแบบการขนานอื่น ๆ เช่น Tensor Parallelism (TP) และ Data Parallelism (DP) ซึ่งสามารถใช้งานได้อย่างง่ายดายโดยใช้ไลบรารีการสื่อสารแบบรวมกลุ่ม (collective communication libraries) เช่น torch.distributed และ NCCL คือ การกำหนดเส้นทางของ MoE จะเกี่ยวข้องกับการสื่อสารแบบเพียร์ทูเพียร์แบบเบาบาง เคอร์เนลการกระจาย (dispatch kernel) จะต้องแบ่งชุดของโทเค็นที่เข้ามาและกระจายไปยังเรดก์ที่โฮสต์ผู้เชี่ยวชาญ ในขณะที่ปัญหาคู่กันอย่างเคอร์เนลการรวม (combine kernel) จะต้องรับโทเค็นที่ประมวลผลแล้วและคำนวณค่าเฉลี่ยถ่วงน้ำหนัก การถ่ายโอนเหล่านี้ไม่แมปกับการสื่อสารแบบรวมกลุ่มและจำเป็นต้องใช้เคอร์เนลที่ปรับแต่งเพื่อให้ได้เวลาแฝงต่ำ

ที่ Perplexity เราได้พัฒนาทั้งเคอร์เนลข้ามโหนดแบบพกพาและเคอร์เนลภายในโหนดเฉพาะทาง โดยนำเสนอผลงานของเราในเชิงลึกยิ่งขึ้นในบทความก่อนหน้านี้เกี่ยวกับ การสื่อสาร MoE เรายังแสดงให้เห็นในบล็อกโพสต์อื่นด้วยว่า การปรับใช้งานแบบหลายโหนด สามารถช่วยลดเวลาแฝงและเพิ่มปริมาณงาน (throughput) สำหรับโมเดล MoE ขนาดใหญ่ได้

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

การอนุมานผ่าน InfiniBand

Tensor parallelism และ expert parallelism สามารถปรับขนาดได้อย่างมีประสิทธิภาพภายในโหนดเดียว เนื่องจาก NVLink มีปริมาณงานสูง (900GB/s) และมีเวลาแฝงระดับไมโครวินาที แม้ว่าระบบ NVIDIA GB200 และ GB300 จะมีหน่วยความจำเพียงพอที่จะเก็บน้ำหนักและ KV caches ของโมเดลที่มีพารามิเตอร์ระดับล้านล้าน แต่การปรับใช้งานบนอินสแตนซ์ AWS p5en ที่มี GPU H200 สูงสุด 8 ตัว จะถูกจำกัดไว้ที่ HBM ขนาด 1120GB ซึ่งจะต้องแบ่งใช้ระหว่างน้ำหนักโมเดลและ KV caches ดังนั้น โมเดลบางรุ่นจึงจำเป็นต้องมีการปรับใช้งานข้ามโหนดผ่าน InfiniBand

แม้ว่า InfiniBand จะยังคงมีความเร็วสูงมาก แต่ก็มีปริมาณงานเพียง 400Gbps (50GB/s) เท่านั้น ซึ่งเพิ่มเวลาแฝงเข้าไปอีกหลักสิบถึงหลักร้อยไมโครวินาทีสำหรับการกำหนดเส้นทาง MoE โชคดีที่เราสามารถซ่อนเวลาแฝงที่เพิ่มขึ้นเหล่านี้ได้ผ่านการใช้ผู้เชี่ยวชาญแบบแชร์ (shared experts) การประมวลผลแบตช์ขนาดเล็ก (micro-batching) และการซ้อนทับกันระหว่างการประมวลผลและการสื่อสาร เพื่อส่งมอบปริมาณงานที่ใช้งานได้จริงในต้นทุนที่แข่งขันได้ DeepSeek เป็นผู้บุกเบิกเทคนิคเหล่านี้ผ่าน DeepSeek-R1 ซึ่งปรับใช้งานบนคลัสเตอร์ของ GPU H800 และอะแดปเตอร์ NVIDIA ConnectX-7 InfiniBand

บนคลัสเตอร์ AWS สถานการณ์จะแตกต่างออกไปเนื่องจากการใช้ Elastic Fabric Adapters (EFA) แบบกำหนดเอง แม้ว่าสิ่งเหล่านี้จะให้ปริมาณงานสูงสุดถึง 400Gbps ในการสื่อสารแบบรวมกลุ่มที่มักใช้ในเวิร์กโหลดการฝึก แต่ก็ยังหย่อนกว่า ConnectX-7 เล็กน้อยในด้านขนาดข้อความที่แลกเปลี่ยนกันระหว่างการกระจายและการรวมของ MoE นอกจากนี้ EFA ไม่รองรับ GPUDirect Async จึงต้องอาศัยเธรดพร็อกซีของ CPU ที่ทำหน้าที่เชื่อมโยงระหว่าง GPU กับ NIC เพื่อเริ่มการถ่ายโอน นอกเหนือจากความซับซ้อนที่เพิ่มขึ้นแล้ว ธุรกรรม PCIe ที่เพิ่มขึ้นยังสร้างค่าใช้จ่ายแฝง (overhead) ระดับไมโครวินาทีให้กับการถ่ายโอนทั้งหมด เฟรมเวิร์กอย่าง NVSHMEM อนุญาตให้ผู้เขียนเคอร์เนลจัดเตรียมการใช้งานที่เป็นอิสระจากอุปกรณ์ได้ แต่ API ทั่วไปเหล่านี้กลับก่อให้เกิดค่าใช้จ่ายแฝงซึ่งทำให้การปรับใช้งานไม่มีประสิทธิภาพทั้งในแง่ของต้นทุนและเวลาแฝง

pplx-kernels ที่อิงตาม NVSHMEM

เคอร์เนลรุ่นก่อนหน้าของเราได้ทำการสรุป (abstract) การใช้งานการขนส่งเบื้องหลังผ่าน NVSHMEM โดยอาศัย nvshmemx_putmem_signal_nbi_warp ในการถ่ายโอนโทเค็นทีละรายการ โดยใช้ตัวนับอะตอมมิกเพื่อซิงโครไนซ์เพียร์ แม้ว่าวิธีนี้จะให้ประสิทธิภาพที่เหมาะสมบน NIC ConnectX-7 ที่มี IBGDA แต่การใช้งาน ConnectX-7 IBRC ที่ใช้พร็อกซีนั้นช้ากว่าอย่างเห็นได้ชัด ในขณะที่การใช้งาน EFA ไม่สามารถทำเวลาแฝงให้อยู่ต่ำกว่าหนึ่งมิลลิวินาทีได้

เราใช้เคอร์เนลเหล่านี้เป็นสภาพแวดล้อมทดสอบสำหรับฮาร์ดแวร์เบื้องหลัง พร้อมทั้งระบุอินเทอร์เฟซที่เหมาะสมที่สุดเพื่อเปิด ใช้งานกับโมเดลของเรา เราตั้งสมมติฐานว่าการถ่ายโอนและการโต้ตอบที่เกิดขึ้นบ่อยครั้งระหว่างเคอร์เนลกับ NIC หรือเธรดพร็อกซีเป็นตัวขัดขวางประสิทธิภาพ นอกจากนี้ การใช้งานของเราไม่ได้อาศัยการรับประกันใดๆ จากฮาร์ดแวร์นอกเหนือจากความน่าเชื่อถือ ซึ่งบ่งชี้ว่าเราสามารถเร่งความเร็วให้ 100% ขึ้นได้ด้วยการลดค่าใช้จ่ายแฝงของพริมิทิฟการซิงโครไนซ์

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

MoE all-to-all บน EFA

ขยาย

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

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

ทั้งเคอร์เนลการกระจายและการรวมจะถูกแบ่งออกเป็นส่วนผู้ส่งและผู้รับครึ่งหนึ่ง เพื่ออำนวยความสะดวกในการซ้อนทับและการประมวลผลแบตช์ขนาดเล็ก ผู้ส่งจะสลับโทเค็นลงในบัฟเฟอร์ส่ง เพื่อให้สามารถเขียนไปยังเพียร์ผ่านการดำเนินการเขียน RDMA เพียงครั้งเดียว หลังจากดำเนินการส่งแล้ว GPU สามารถทำงานอื่นๆ (เช่น การคูณผู้เชี่ยวชาญแบบแชร์) ในขณะที่การถ่ายโอนผ่านเครือข่ายกำลังรอดำเนินการอยู่เบื้องหลัง เคอร์เนลผู้รับจะบล็อกการทำงานจนกว่าจะได้รับการถ่ายโอนทั้งหมด เบื้องหลังจะมีเธรดงานฝั่งโฮสต์ที่จัดการการถ่ายโอน โดยจะได้รับสัญญาณจาก GPU เมื่อบัฟเฟอร์ส่งถูกเติมข้อมูลแล้ว หลังจากนั้นจะทำการกระจายการเขียน RDMA เมื่อ NIC ส่งสัญญาณเสร็จสิ้น เธรดจะปลดบล็อกผู้รับ ซึ่งจะส่งโทเค็นกลับไปยังโมเดล

เคอร์เนลมีการเชื่อมโยงอย่างแนบแน่นกับโฮสต์ โดยสื่อสารกันผ่านทั้ง unified memory และ GDRCopy CUDA unified memory ช่วยให้สามารถแมปหน่วยความจำของอุปกรณ์หรือหน่วยความจำของโฮสต์เข้าไปใน address space ของทั้งสองฝั่งได้ โดยจัดการการสื่อสารผ่าน PCIe อย่างโปร่งใส

แม้ว่า unified memory จะเหมาะสมกับการถ่ายโอนจำนวนมาก แต่การสำรวจสถานะด้วยเวลาแฝงต่ำ (low-latency polling) จะทำผ่าน GDRCopy โดย GDRCopy ใช้ GPUDirect RDMA เพื่ออ่านและเขียนหน่วยความจำ GPU ด้วยเวลาแฝงระดับไมโครวินาที

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

บัฟเฟอร์ส่งและรับ (Send and Receive Buffers)

การออกแบบบัฟเฟอร์ส่งสำหรับเคอร์เนลการกระจายถือเป็นสิ่งสำคัญอย่างยิ่งในการบรรลุเวลาแฝงที่ใช้งานได้จริงเมื่อต้องเกี่ยวข้องกับพร็อกซีโฮสต์ หากมีเรดก์ N เรดก์ โดยแต่ละเรดก์กระจายโทเค็นจำนวน T โทเค็น ซึ่งแต่ละโทเค็นกำหนดเส้นไปยังผู้เชี่ยวชาญจำนวน R ตัวจากทั้งหมด E ตัว เรดก์หนึ่งเรดก์จะสามารถรับโทเค็นได้สูงสุดถึง N * T * max(R, E / N) โทเค็นในกรณีที่แย่ที่สุดที่เป็นไปได้ หากโทเค็นทั้งหมดถูกกำหนดเส้นทางมาที่เรดก์นั้น นอกจากนี้ ขอบเขตขั้นต่ำนี้จะสามารถบรรลุได้ก็ต่อเมื่อผู้ส่งจัดเรียงการเขียนของตนแบบต่อเนื่องกัน (contiguously) เนื่องจากบัฟเฟอร์รับแบบส่วนตัวจะใช้หน่วยความจำมากเกินไป เพื่อให้ผู้ส่งสามารถจัดเรียงการเขียนติดต่อกันได้โดยไม่มีความขัดแย้ง จำเป็นต้องมีการกระจายข้อมูลการกำหนดเส้นทางทั้งหมดไปยังแต่ละเรดก์ ดังนั้น เคอร์เนลการกระจายของเราจึงทำการแลกเปลี่ยนจำนวนโทเค็นต่อผู้เชี่ยวชาญก่อน หลังจากนั้นแต่ละเคอร์เนลจึงจะสามารถระบุได้ว่าการเขียนของตนจะไปที่ใดและได้รับโทเค็นจากเรดก์ใดที่ออฟเซ็ตไหน

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

ขยาย

เคอร์เนลการรวมจะมีความเรียบง่ายกว่ามาก เนื่องจากข้อมูลการกำหนดเส้นทางมีอยู่แล้วจากขั้นตอนการกระจาย เคอร์เนลสามารถคัดลอกโทเค็นผ่านการเขียนครั้งเดียวต่อเพียร์หนึ่งตัว หลังจากนั้นจะอ่านจากบัฟเฟอร์รับเพื่อคำนวณค่าเฉลี่ยถ่วงน้ำหนักของโทเค็นลงในเทนเซอร์ผลลัพธ์

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

การกระจายและการรวม (Dispatch and Combine)

ภาพประกอบด้านล่างนี้ให้รายละเอียดเกี่ยวกับเวลาและการโต้ตอบระหว่าง GPU, CPU และ NIC โดยแสดงการวนซ้ำหนึ่งครั้งของการกระจายและการรวมที่สลับกับการทำงานของเคอร์เนลของโมเดล

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

ขยาย

ขยาย

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

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

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

TransferEngine สำหรับการกำหนดเส้นทาง MoE

ในขั้นต้นเพื่อสร้างต้นแบบการใช้งาน เราได้นำ Transfer Engine ที่สร้างขึ้นสำหรับการพรีฟิลและการถอดรหัสแบบแยกส่วน (disaggregate prefill and decode) กลับมาใช้ใหม่เพื่อออกคำสั่งและติดตามการเขียน RDMA อย่างไรก็ตาม ไม่เหมือนกับการถอดรหัสแบบแยกส่วนที่แต่ละ node โพสต์การเขียนจำนวนมากไปยังเพียร์เดียวกัน การกระจายของ MoE จะโพสต์การถ่ายโอนขนาดเล็กไปยังเพียร์ทั้งหมดในครั้งเดียว การทำงานแบบ out-of-the-box นั้นให้ประสิทธิภาพที่ไม่เหมาะสมนัก และการใช้งานของเราก็ตามหลัง DeepEP บน ConnectX-7 อยู่ประมาณ 20us ดังนั้นเราจึงปรับแต่ง TransferEngine ให้เชี่ยวชาญเพื่อรองรับการทำงานสองอย่าง ได้แก่ scatter และ barrier เพื่อเพิ่มประสิทธิภาพการโต้ตอบกับเพียร์จำนวนมาก

ทั้งการทำงานแบบ scatter และ barrier จะเปิดเผยกลุ่มเพียร์ที่สามารถลงทะเบียนล่วงหน้าเพื่อเติมข้อมูลล่วงหน้าให้กับโครงสร้างข้อมูลที่ใช้โดยการใช้งานการขนส่งเบื้องหลัง (การทำเทมเพลต Work Request (WR)) Scatter จะคัดลอกชิ้นส่วนที่แตกต่างกันจากบัฟเฟอร์ต้นทางเดียวกันไปยังออฟเซ็ตที่แตกต่างกันบนสมาชิกทั้งหมดของกลุ่มเพียร์ ส่วน Barrier จะส่งสัญญาณไปยังเพียร์ทั้งหมดผ่านสัญญาณทันที การจัดคิวและการจัดการการทำงานเหล่านี้แบบกลุ่มจะช่วยประหยัดไมโครวินาทีอันมีค่าในเส้นทางวิกฤต บน EFA ซึ่งมี NIC 200Gbps สองตัวรวมกันเพื่อให้แบนด์วิดท์ 400Gbps เราจะทำการแบ่งชาร์ดกลุ่มเพียร์ข้าม NIC ทั้งสองตัว แทนที่จะแบ่งชาร์ดไบต์ของการถ่ายโอนแต่ละครั้ง

แม้ว่าในตอนแรกเราจะพัฒนา TransferEngine บน EFA โดยใช้ libfabric เท่านั้น แต่เคอร์เนล MoE ก็สามารถทำงานบน ConnectX-7 ได้โดยไม่ต้องเปลี่ยนโค้ดเลยแม้แต่บรรทัดเดียว หลังจากเพิ่มการรองรับ ConnectX-7 ใน TransferEngine โดยใช้ libibverbs เมื่อเปรียบเทียบกับ EFA ที่อิงตาม Scalable Reliable Datagram (SRD) แล้ว ConnectX-7 จำเป็นต้องมีงานเพิ่มเติมสำหรับการตั้งค่าการเชื่อมต่อและการจัดการเพียร์ เราใช้คู่คิว Unreliable Datagram (UD) เพื่อบูตสแตรป Reliable Connection (RC) เราสร้างคู่คิวสองคู่สำหรับเพียร์ระยะไกลแต่ละตัว คู่หนึ่งสำหรับ SEND/RECV สองทาง และอีกคู่สำหรับ WRITE_IMM ทางเดียว เนื่องจากทั้ง RECV และ WRITE_IMM ต่างก็ใช้การดำเนินการ RECV ตามลำดับ ในขณะที่การดำเนินการทั้งสองแบบต้องการขนาดบัฟเฟอร์ที่แตกต่างกัน นอกจากนี้ เรายังผ่อนปรนลำดับธุรกรรม PCIe ระหว่าง ConnectX-7 และหน่วยความจำ GPU ด้วย เช่นเดียวกับ EFA เราดำเนินการปรับแต่งการทำเทมเพลต WR นอกจากนี้ เมื่อส่งข้อความขนาดเล็กหลายข้อความไปยังเพียร์เดียวกัน เราจะใช้การเชื่อมโยง WR (WR chaining) เพื่อปรับปรุงอัตราข้อความ

เนื่องจากการทำงานฝั่ง GPU เหมือนกันในอะแดปเตอร์ทั้งสองแบบ เราจึงผลักดันการใช้งานการขนส่งเบื้องหลังให้ถึงขีดจำกัดเพื่อเหนือกว่าประสิทธิภาพของ DeepEP ซึ่งเป็นการเน้นย้ำว่าช่องว่างระหว่างการกำหนดเส้นทาง MoE ของ EFA และ ConnectX-7 นั้นถูกกำหนดโดยความสามารถของฮาร์ดแวร์เบื้องหลัง

การประเมินผล (Evaluation)

เวลาแฝงของการกระจาย/การรวม (Dispatch/Combine Latency)

เราได้ประเมินเวลาแฝงของการกระจายและการรวมบนระบบ H200 ที่ติดตั้งอะแดปเตอร์ ConnectX-7 และ EFA บนการกำหนดค่ามาตรฐานของ DeepSeek-V3 และ Kimi-K2 โดยมีมิติที่ซ่อนอยู่ (hidden dimension) ที่ 7168, การกระจายแบบ fp8 ที่ปรับขนาดตามบล็อก (block-scaled) และการรวมแบบ bf16 เราพิจารณาขนาดแบตช์ที่ 128 สำหรับการกระจายและ 4096 สำหรับการรวม บน ConnectX-7 เราเปรียบเทียบกับ DeepEP และ pplx-kernels ที่อิงตาม NVSHMEM รุ่นก่อนหน้าของเรา โดยใช้ทั้ง GPUDirect Async (IBGDA) และพร็อกซีโฮสต์ (IBRC) เราบรรลุเวลาแฝงรวมที่ 459us, 582us และ 692us บน EP16, EP32 และ EP64 ตามลำดับ เรามีประสิทธิภาพเหนือกว่าเคอร์เนลที่อิงตาม EFA ตัวอื่นๆ ที่อาศัยพร็อกซีทั่วไป เช่น UCCL-EP ซึ่งทำเวลาแฝงได้ 519us, 966us และ 1159us

ขยาย

ในการถอดรหัส แม้ว่าเคอร์เนลของเราจะช้ากว่า DeepEP อยู่ไม่กี่ไมโครวินาที แต่เราก็บรรลุเวลาแฝงโดยรวมที่ต่ำกว่า เนื่องจากการรวมเคอร์เนลของเรามีความเร็วกว่าอย่างมาก แม้ว่าเราจะใช้พร็อกซีโฮสต์ แต่เนื่องจากการถ่ายโอนข้อมูลจำนวนมากของเรา ทำให้เรามีประสิทธิภาพเหนือกว่าการใช้งานแบบไม่มีพร็อกซีที่อิงตาม IBGDA ทั้งหมด เรายังแสดงให้เห็นด้วยว่าการใช้งานพร็อกซีมีความสำคัญ เนื่องจากพร็อกซีที่ได้รับการปรับแต่งมาเป็นพิเศษของเรามีความเร็วกว่าการใช้งาน NVSHMEM ทั่วไปถึงหนึ่งลำดับความสำคัญ แม้ว่า EFA จะช้ากว่า ConnectX-7 ถึงสองเท่าเมื่อถ่ายโอนแพ็กเก็ตขนาด 256KB ซึ่งใกล้เคียงกับขนาดแพ็กเก็ตทั่วไปที่แลกเปลี่ยนโดยการกำหนดเส้นทาง MoE แต่ช่องว่างระหว่างอะแดปเตอร์ทั้งสองโดยรวมแล้วก็น้อยกว่าอย่างเห็นได้ชัด

ขยาย

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

ปริมาณงาน (Throughput) ของ Kimi-K2 / DeepSeek-V3

เพื่อพิสูจน์ความสามารถในการใช้งานจริงของเคอร์เนลของเรา เราได้ทำการทดสอบประสิทธิภาพแบบ end-to-end ของเอ็นจิ้นการอนุมานภายในองค์กรของเราที่ให้บริการ DeepSeek-V3 (พารามิเตอร์ 671B) และ Kimi-K2 (พารามิเตอร์ 1T) บนอินสแตนซ์ p5en หลายตัวที่ติดตั้ง Elastic Fabric Adapters (EFA) และ GPU H200 เราเรียกใช้โมเดลในโหมดตรวจสอบร่าง (draft verification mode) และรายงานเวลาแฝงโดยสมมติว่าอัตราการยอมรับอยู่ที่ 80% โดยใช้เลเยอร์ MTP เดี่ยว

บน DeepSeek-V3 เราทำการเปรียบเทียบกับเบสไลน์ภายในโหนด EP=8 ที่ใช้งานเคอร์เนลการกระจายและการรวมเฉพาะ NVLink ที่มีประสิทธิภาพ การปรับใช้งานข้ามโหนดโดยรวมแล้วมีประสิทธิภาพเทียบเท่าหรือเหนือกว่าการปรับใช้งานแบบโหนดเดี่ยว โดยเฉพาะอย่างยิ่งในแบตช์ขนาดกลาง ซึ่งแสดงให้เห็นว่าการปรับขนาดข้ามโหนดจะนำมาซึ่งการปรับปรุงปริมาณงานโดยรวมอย่างแท้จริง

บน Kimi-K2 เคอร์เนลเหล่านี้ทำให้สามารถให้บริการได้ เนื่องจากขนาดของโมเดล ทำให้ไม่สามารถใส่ลงในโหนดเดียวกันโดยใช้ Data Parallelism (DP) บนโหนด H200 ได้ ดังนั้นจึงจำเป็นต้องมีการปรับใช้งานแบบข้ามโหนดเพื่อให้บริการโมเดลนี้ เนื่องจากโมเดลมีฮีดแอตเทนชัน (attention heads) น้อยกว่า เราจึงบรรลุเวลาแฝงที่ดีกว่า DeepSeek-V3 และเราสามารถให้บริการโมเดลบน EFA ด้วยเวลาแฝงที่ใช้งานได้จริง

ขยาย

ขยาย

เรายังได้พิจารณาการประมวลผลแบตช์ขนาดเล็ก (micro-batching) ด้วย แม้ว่าจะช่วยปรับปรุงเวลาแฝงได้เฉพาะเมื่อขนาดแบตช์สูงมากๆ หรือเมื่อใช้โทเค็นร่างมากกว่าหนึ่งโทเค็นก็ตาม เพื่อให้ได้รับประโยชน์จากการแบ่งปัญหากลุ่ม GEMM ออกเป็นสองส่วน ขนาดแบตช์ของโทเค็นจะต้องมีความหนาแน่นเพียงพอ ส่วนใหญ่ของเวลายังคงถูกใช้ไปกับการกระจายและการรวมของ MoE ซึ่งบ่งชี้ว่าการปรับปรุงใดๆ เพิ่มเติมต่อเคอร์เนลจะช่วยปรับปรุงเวลาแฝงและอำนวยความสะดวกในการทำ micro-batching ตามไปด้วย

งานในอนาคต (Future Work)

เรากำลังทำงานอย่างใกล้ชิดกับวิศวกรของ AWS เพื่อปรับปรุงประสิทธิภาพบน EFA ในส่วนของเรา เราจะติดตามการอัปเดต libfabric ซึ่งจะช่วยลดค่าใช้จ่ายแฝงของระนาบข้อมูล (data plane overhead) นอกจากนี้ เราจะทำการทดลองกับ efa-direct เพื่อลดค่าใช้จ่ายแฝงของสเปซผู้ใช้ (userspace overheads) ให้ดียิ่งขึ้น และเชื่อมต่อกับฮาร์ดแวร์โดยตรง

ข้อมูลอ้างอิง (References)

รายงานทางเทคนิค DeepSeek R1

รายงานทางเทคนิค Kimi K2

DeepEP

UCCL-EP

การปรับปรุงประสิทธิภาพเครือข่ายของระบบ HPC โดยใช้ NVIDIA Magnum IO NVSHMEM และ GPUDirect Async

การอ้างอิง

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}