CuTeDSL ที่ Perplexity

วิธีที่เราใช้ CuTeDSL เพื่อสร้าง GPU เคอร์เนลที่นำโมเดลของเราไปสู่ประสิทธิภาพสูงสุดบน GPU NVIDIA

ผู้เขียนPerplexity Research

ที่ Perplexity เราใช้งานเอ็นจิ้น Runtime-Optimized Serving Engine (ROSE) แบบอินเฮาส์เพื่อให้บริการโมเดลตั้งแต่เอ็มเบดดิ้งไปจนถึงโมเดลภาษาขนาดใหญ่ (LLM) แบบพารามิเตอร์นับล้านล้านตัว เอ็นจิ้นของเราซึ่งอยู่เบื้องหลัง API ทั้งหมด รวมถึง Sonar, Search และ Embeddings ถูกสร้างขึ้นเพื่อให้เหล่านักวิจัยสามารถสำรวจโมเดลและสถาปัตยกรรมที่เหมาะสมที่สุด พร้อมทั้งปรับใช้โมเดลที่มีแนวโน้มดีที่สุดเข้าสู่ระบบจริงได้อย่างรวดเร็ว เพื่อให้บรรลุเป้าหมายนี้ เราจึงจำเป็นต้องมีความสามารถในการใช้งาน ปรับใช้ และเพิ่มประสิทธิภาพการอนุมานสำหรับสถาปัตยกรรมล้ำสมัยได้อย่างรวดเร็ว ในบล็อกโพสต์นี้ เราจะมาสำรวจว่าเคอร์เนล CuTeDSL ที่ผสานรวมเข้ากับเอ็นจิ้นการอนุมานของเราช่วยให้เราสามารถสร้างเคอร์เนล GPU ที่จำเป็น เพื่อยำโมเดลไปสู่ประสิทธิภาพสูงสุดบน GPU ของ NVIDIA ได้อย่างไร

การอนุมานที่ Perplexity AI

โมเดลแบบกำหนดเองที่สร้างโดยทีมต่างๆ ที่ Perplexity โฮสต์แบบอินเฮาส์บน GPU NVIDIA Hopper และ Blackwell โดยใช้เอ็นจิ้นการอนุมาน ROSE ที่สร้างโดยทีม AI Inference เริ่มแรก ROSE ถูกสร้างขึ้นเพื่อให้บริการโมเดล Llama ที่ปรับแต่งสำหรับทั้งการถอดรหัสภาษาและการจัดหมวดหมู่ได้ดียิ่งขึ้น ด้วยอินเทอร์เฟซที่เข้ากันได้กับ NVIDIA Triton Inference Server เมื่อเวลาผ่านไป เอ็นจิ้นนี้ได้เติบโตจนกลายเป็นเอ็นจิ้นการอนุมานเต็มรูปแบบที่ให้บริการ LLM ขนาดใหญ่ รวมถึงโมเดลสำหรับการจัดอันดับ การจัดหมวดหมู่ การให้คะแนน และโมเดลอื่นๆ ที่ใช้ทรานส์ฟอร์เมอร์สำหรับทีมหลากหลายกลุ่ม เพื่อขับเคลื่อนผลิตภัณฑ์ต่างๆ ของ Perplexity เป้าหมายการออกแบบของ ROSE คือความเรียบง่าย: เมื่อเวลาผ่านไป มีการลงทุนลงแรงเพื่อให้ได้ประสิทธิภาพสูงสุด แต่ลำดับความสำคัญหลักคือการปรับตัวให้เข้ากับสถาปัตยกรรมและกรณีการใช้งานที่พัฒนาอยู่เสมอสำหรับโมเดลของเรา

พื้นที่การกำหนดค่าของเคอร์เนลครอบคลุมประเภทข้อมูล มิติของเฮด ค่าคงที่ขณะคอมไพล์ และการตั้งค่าการเปิดใช้งาน

ในระดับสูง ROSE ประกอบด้วย เอ็นจิ้น ที่ทำหน้าที่ปรับโมเดลให้เข้ากับอินเทอร์เฟซที่ไคลเอนต์สามารถโต้ตอบได้ เอ็นจิ้นเหล่านี้มีความสามารถในการนำกลับมาใช้ใหม่สูง โดยจัดการการจัดกำหนดการและการจัดกลุ่มคำขอ (request scheduling and batching) รวมถึงดึงข้อมูลโทเค็นที่ถอดรหัสแล้ว เอ็มเบดดิ้ง หรือคะแนนจากโมเดลเบื้องหลัง การเริ่มต้นทำงานของอุปกรณ์ ช่องทางการสื่อสารระหว่างกระบวนการและระหว่างโหนด และการโหลดน้ำหนัก (weight loading) ล้วนได้รับการจัดการโดยเอ็นจิ้นเหล่านี้ สำหรับ LLM เอ็นจิ้นมีหน้าที่รับผิดชอบในการจัดกลุ่ม (batching) การแบ่งส่วน (chunking) และการสุ่มตัวอย่าง (sampling) นอกจากนี้ยังมีหน้าที่ในการจัดสรรพื้นที่เก็บข้อมูล KV ทั้งสำหรับ attention แบบเต็มและแบบเชิงเส้น รวมถึงการจับคู่คำนำหน้า (prefix matching) เอ็นจิ้นเอ็มเบดดิ้งจะมีความเรียบง่ายกว่า เนื่องจากจะทำการจัดกลุ่มแบบออนไลน์เท่านั้น การดำเนินการได้รับการปรับ optimization ให้เอ็นจิ้นเตรียมงานรอบถัดไปในขณะที่โมเดลกำลังทำงานบนตัวเร่งความเร็วไปพร้อมๆ กัน

ROSE เปิดเผยชุดเลเยอร์แบบกำหนดเองที่มีความอุดมสมบูรณ์สำหรับการสร้างโมเดล โดยห่อหุ้ม GPU เคอร์เนลที่สร้างด้วย CuTeDSL, Triton, CUDA, CUTLASS และ cuBLAS สำหรับการคูณเมทริกซ์และการใช้งาน attention เราพึ่งพาเคอร์เนล NVIDIA จาก CUTLASS และ cuBLAS เป็นหลัก ส่วนที่เหลือของไปพไลน์ รวมถึงเอ็มเบดดิ้ง, norm, MoE และเคอร์เนลแอคติเวชัน จะเป็นส่วนหนึ่งของไลบรารีเคอร์เนลของ ROSE เพื่อรองรับฮาร์ดแวร์ที่เราใช้รันการอนุมานได้ดีที่สุด ROSE จึงติดตั้ง primitive การสื่อสารระหว่างอุปกรณ์ของตัวเองสำหรับการทำ MoE routing, dispatch และ combine แม้ว่าเคอร์เนลอื่นๆ จะค่อนข้างเรียบง่าย แต่ก็ต้องได้รับการปรับแต่งอย่างละเอียดเพื่อให้ได้ประสิทธิภาพสูงสุดในมิติซ่อนเร้นและอุปกรณ์ต่างๆ ทั้งหมดที่รันอยู่ นอกจากนี้ เรายังพึ่งพาการปรับแต่งแบบคอมไพล์ไทม์เพื่อรองรับพารามิเตอร์และการปรับแต่งทั้งหมดของโมเดลต่างๆ (เช่น weight bias เทียบกับไม่มี weight bias สำหรับ RMS norm, ฟังก์ชันแอคติเวชัน ฯลฯ) ดังนั้นเราจึงตัดสินใจใช้ DSL เพื่อใช้งานสิ่งเหล่านี้: โดยเริ่มแรกเราสร้าง ROSE โดยใช้ Triton อย่างไรก็ตาม นับแต่นั้นมาเราได้นำ CuTeDSL มาใช้เป็นสภาพแวดล้อมหลักสำหรับการเขียนโปรแกรม GPU

ทำไมต้อง CuTeDSL?

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

CuTeDSL เป็น DSL ที่ใช้ Python สำหรับ GPU ของ NVIDIA ซึ่งสร้างขึ้นบนพื้นฐานของพีชคณิตเลย์เอาต์ CuTe และ MLIR โดยจะคอมไพล์แบบ Just-in-Time ไปเป็น PTX ที่ได้รับการปรับ optimization อย่างสูง ในขณะที่ยังคงรักษาการควบคุม primitive ฮาร์ดแวร์ระดับต่ำไว้ได้ เป้าหมายคือการมอบฟีเจอร์และประสิทธิภาพที่เทียบเท่ากับ CUTLASS C++ โดยปราศจากเวลาในการคอมไพล์ที่ยาวนานซึ่งขัดขวางความเร็วในการพัฒนา จากภาษาที่เราได้สำรวจหรือใช้งานมา (รวมถึง CUDA, cuTile, TVM IR และ Triton) เราพิจารณาว่าภาษานี้มีความเหมาะสมที่สุดสำหรับการพัฒนาเคอร์เนลของเรา

โดยรวมแล้ว แม้ว่าจะไม่ได้อยู่ในระดับสูงและกระชับเท่า Triton หรือ cuTile แต่เราพบว่า CuTe Layout Algebra และองค์ประกอบต่างๆ ของมัน (ซึ่งเปิดเผยผ่าน cute.Tensor) เป็นตัวแทนที่สมเหตุสมผล นอกจากนี้ CuTeDSL ยังให้การเข้าถึงฮาร์ดแวร์ primitive แบบไม่มีข้อจำกัด ทำให้เราสามารถควบคุมในระดับที่ต้องการเพื่อให้บรรลุตะวันประสิทธิภาพสูงสุดในทุก edge cases สำหรับ CUDA โอกาสในการปรับแต่งแบบคอมไพล์ไทม์นั้นมีจำกัด เนื่องจากพารามิเตอร์เพิ่มเติมแต่ละตัวเพิ่มต้นทุนที่ไม่สมเหตุสมผลให้กับเวลาในการคอมไพล์หรือความซับซ้อนให้กับระบบบิวด์ ในทางตรงกันข้าม ด้วย DSL ที่คอมไพล์แบบ Just-In-Time เราสามารถปรับแต่งได้อย่างเข้มข้นยิ่งขึ้นโดยไม่ต้องแลกมาด้วยความเร็วในการทำซ้ำ (iteration pace) แตกต่างจาก TVM IR ซึ่งใช้ CUDA เป็นขั้นตอนกลาง เราพบว่าไปพไลน์ที่ใช้ MLIR ซึ่งลดระดับลงไปเป็น PTX นั้นมีความทนทานและง่ายต่อการดีบักมากกว่าไปพไลน์ที่ปล่อย CUDA ออกมาเป็นขั้นตอนกลาง นอกจากนี้ ด้วยการใช้ custom ops เพื่อห่อหุ้มคำสั่ง PTX เราพบว่าเราสามารถใช้ประโยชน์จากฟีเจอร์ทั้งหมดของฮาร์ดแวร์เบื้องหลังได้อย่างเต็มที่

ประสบการณ์นักพัฒนา

ต้นทุนการคอมไพล์ส่วนใหญ่ของเคอร์เนล CUDA มักจะหมดไปกับขั้นตอนการขยายเทมเพลต (template expansion phase) ของเคอร์เนล เนื่องจากการ optimization และการคอมไพล์ตั้งแต่ขั้นตอน PTX เป็นต้นไปสามารถเสร็จสิ้นได้ภายในเวลาไม่กี่วินาที ด้วยพารามิเตอร์สำหรับคอมไพล์ไทม์จำนวนมาก ขั้นตอนนี้ของการคอมไพล์อาจพุ่งสูงขึ้นเป็นเวลาหลายนาทีสำหรับเคอร์เนลที่ไม่ใช่เรื่องง่าย คอขวดนี้ยิ่งชัดเจนขึ้นจากความจริงที่ว่าโอกาสในการคอมไพล์แบบคู่ขนานจะเกิดขึ้นทีหลังเท่านั้น ในระดับ IR และ PTX โดยที่การจัดการเทมเพลต C++ จะเกิดขึ้นทีละขั้นตอนตามลำดับ แม้ว่าจะมีเทคนิคต่างๆ เพื่อแก้ไขปัญหานี้ เช่น การย้ายการสร้างเทมเพลต (template instantiations) ไปยังยูนิตการคอมไพล์แยกต่างหาก แต่สิ่งเหล่านี้จำเป็นต้องมีการปรับแต่งระบบบิวด์อย่างมากและการใช้แมโครที่ยุ่งยาก ในทางตรงกันข้าม เคอร์เนล CuTeDSL ที่คอมไพล์แบบ Just-In-Time (JIT) ไม่จำเป็นต้องมีกฎการบิวด์ที่ซับซ้อนและไม่เป็นอุปสรรคต่อเวลาในการคอมไพล์ เนื่องจากวงจรการพัฒนาสามารถมุ่งเน้นไปที่จุดเดียวในสเปซการกำหนดค่า ในขณะที่การทดสอบยูนิตเทสต์อย่างระมัดระวังในภายหลังสามารถครอบคลุมสเปซการออกแบบทั้งหมดของเคอร์เนลได้

แม้ว่าการคอมไพล์แบบ JIT จะย้ายต้นทุนจากช่วงพัฒนาไปยังช่วงรันไทม์ แต่โอเวอร์เฮดนั้นถือว่ายอมรับได้อย่างมาก เมื่อโมเดลถูกนำไปใช้งาน ระบบจะเปิดใช้งานเพียงบางส่วนจากจุดต่างๆ ทั้งหมดในพื้นที่การกำหนดค่าทั้งหมด ทำให้เกิดโอเวอร์เฮดระดับวินาทีในระหว่างการเริ่มต้น/อุ่นเครื่องของเอนจินอนุมาน ตรงกันข้าม การพรีคอมไพล์การกำหนดค่าทั้งหมดของเคอร์เนลอาจใช้เวลาหลายนาที นอกจากนี้ ยังสามารถลดโอเวอร์เฮดเพิ่มเติมได้ด้วยการแคชอาร์ทิแฟกต์ที่คอมไพล์ด้วย JIT

ต้นทุนการคอมไพล์ที่สูงของ CUDA ยิ่งเด่นชัดขึ้นจากความยากลำบากในการติดตามที่มาของข้อผิดพลาดเมื่อมีการใช้เทมเพลตดิสแพทช์และการปรับแต่ง (template-based dispatch and specialization) เช่นเดียวกับในกรณีของ CuTe และ CUTLASS เมื่อใช้ DSL ข้อผิดพลาดสามารถชี้ไปยังบรรทัดโค้ดที่สร้างมันขึ้นมาได้ ซึ่งช่วยให้การดีบักง่ายขึ้นอย่างมาก นอกจากนี้ การแมปโค้ดที่สร้างขึ้นไปยังคำสั่ง DSL ที่สร้างมันขึ้นมาผ่านขั้นตอน MLIR ระดับกลาง หรือการอาศัยข้อมูลการดีบักที่แนบมากับ PTX ที่สร้างขึ้น ทำได้ง่ายขึ้นมาก ด้วยแง่มุมเหล่านี้ของ DSL เราพบว่าตนเองมีผลผลิตในการพัฒนาเคอร์เนลเพิ่มขึ้นอย่างมาก

สเปซการกำหนดค่าขนาดใหญ่

แม้ว่าเคอร์เนลจะมีส่วนที่คล้ายคลึงกันมากมายในแต่ละโมเดล แต่เคอร์เนลเหล่านี้ก็ยังคงต้องติดตั้งปุ่มปรับการกำหนดค่า (configuration knobs) หลากหลายแบบที่อาจส่งผลต่อประสิทธิภาพการทำงาน พารามิเตอร์ที่เกี่ยวข้องประกอบด้วย:

มิติซ่อนเร้น/มิติเฮด/จำนวนแอตเทนชันเฮด (Hidden Dimension/Head Dimension/Number of Attention Heads)

ประเภทข้อมูลอินพุต/เอาต์พุต (Input/Output data types)

รูปแบบการควอนไทเซชัน (Quantization scheme)

ค่าคงที่ของคอมไพล์ไทม์ (ไบแอส, ฟังก์ชันแอคติเวชัน ฯลฯ)

สถาปัตยกรรม ROSE รวมถึงเอ็นจิน โมเดล ไลบรารีเคอร์เนล และบริการต่างๆ

การปรับแต่งแบบคอมไพล์ไทม์มีประโยชน์ทั้งในด้านประสิทธิภาพและการดีบัก โดยเฉพาะในโปรเซสเซอร์แบบไปพไลน์ เช่นเดียวกับในกรณีของ GPU สาขาที่ไม่มีอยู่จริง (nonexistent branch) ย่อมดีกว่าสาขาที่ถูกเลือกเสมอหรือไม่มีวันถูกเลือกเลย นอกจากนี้ การคอมไพล์เคอร์เนลสำหรับชุดปุ่มปรับการกำหนดค่าเฉพาะจะส่งผลให้มีคำสั่ง PTX/SASS น้อยลง ซึ่งง่ายต่อการวิเคราะห์และดีบักเมื่อใช้ cuda-gdb ระหว่างการตรวจสอบ

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

การกรอกข้อมูลล่วงหน้าและการถอดรหัสความเชี่ยวชาญพิเศษ

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

พรีฟิล (Prefill): ขนาดแบตช์ขนาดเล็ก (จำนวนชังก์และซีเควนซ์น้อย), มีโทเค็นจำนวนมากต่อซีเควนซ์ ต้องเพิ่ม throughput ให้สูงสุดสำหรับทั้งเอ็มเบดดิ้งและอินพุตโทเค็นของ LLM โดยมี Latency อยู่ในช่วงหลักสิบถึงหลักร้อยมิลลิวินาที

การถอดรหัส (Decode): ขนาดแบตช์ตั้งแต่ขนาดเล็กไปจนถึงขนาดใหญ่, โทเค็นหนึ่งตัว (ไม่มีการถอดรหัสแบบคาดเดา) ไปจนถึงหลายโทเค็นต่อซีเควนซ์ (การคาดเดาแบบร่างเป้าหมาย, การทำนายหลายโทเค็น) มี Latency ใกล้เคียงกับ Latency ของการเรียกใช้เคอร์เนล (kernel launch latency)

สำหรับ GEMM และ attention ขั้นตอนพรีฟิลจะทำงานบนโทเค็นจำนวนมากและประสิทธิภาพจะถูกจำกัดด้วย throughput การคำนวณของอุปกรณ์ ส่วนการถอดรหัสจะเกี่ยวข้องกับการดำเนินการเดียวกันกับขนาดแบตช์ขนาดเล็ก โดยเวลาในการทำงานจะถูกครอบงำด้วยต้นทุนในการโหลดแอคติเวชันและน้ำหนักจาก DRAM เข้าสู่สตรีมมิงมัลติโปรเซสเซอร์ (streaming multiprocessors) และเทนเซอร์คอร์ ซึ่งถูกจำกัดด้วยแบนด์วิดท์หน่วยความจำของอุปกรณ์แทน แม้ว่าขนาดแบตช์หรือจำนวนโทเค็นทั้งหมดในแบตช์จะเป็นตัวบ่งชี้ที่ค่อนข้างชัดเจน แต่เกณฑ์การสลับระหว่างพรีฟิลและการถอดรหัสมีความละเอียดอ่อนมากกว่า เนื่องจากแบนด์วิดท์หน่วยความจำถูกกำหนดโดยมิติซ่อนเร้นของโมเดลและประเภทข้อมูลของแอคติเวชันด้วย ดังนั้น เคอร์เนลจึงถูกแยกออกเป็นการใช้งานแบบพรีฟิลและการถอดรหัส หากความแตกต่างทางอัลกอริทึมหรือการกำหนดกริด/บล็อกก่อให้เกิดข้อได้เปรียบแก่ฝ่ายใดฝ่ายหนึ่ง

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

cpp
template <size_t NUM_WARPS> __global__ void __launch_bounds__(NUM_WARPS * WARP_SIZE, 1) some_kernel(...) { }

จำนวนรีจิสเตอร์ที่มีให้สำหรับเธรดทั้งหมดในบล็อกคือข้อจำกัดของฮาร์ดแวร์ ซึ่งเท่ากับ 64K บน Hopper และ Blackwell หากทราบจำนวนเธรดล่วงหน้า คอมไพล์เลอร์จะทราบแน่ชัดว่าสามารถกำหนดรีจิสเตอร์ให้แต่ละเธรดได้จำนวนเท่าใด และสามารถประสบความสำเร็จในการคอมไพล์เคอร์เนลที่ได้รับการปรับ optimization อย่างเข้มข้นยิ่งขึ้น อย่างไรก็ตาม สิ่งนี้แลกมาด้วยต้นทุนในเวลาในการคอมไพล์ เนื่องจากตอนนี้เคอร์เนลจะต้องถูกคอมไพล์ใหม่สำหรับขนาดวอร์ปที่แตกต่างกัน ด้วย CuTeDSL ปัญหานี้จะไม่เกิดขึ้นอีกต่อไป เนื่องจากเราไม่ต้องแบกรับต้นทุนการคอมไพล์สำหรับสเปซการกำหนดค่าทั้งหมดล่วงหน้า นอกจากนี้ DSL ยังคาดหวังให้ขนาดบล็อกเป็นนิพจน์ค่าคงที่ โดยทำ optimization ที่นัยโดย __latch_bounds__ ใน CUDA โดยปริยาย

นอกเหนือจากการทำให้จำนวนวอร์ปเป็นค่าคงที่แล้ว การปรับเปลี่ยนจำนวนวอร์ประหว่างพรีฟิลและการถอดรหัสยังเป็นสิ่งที่พึงประสงค์ เพื่อให้ใช้งานแบนด์วิดท์หน่วยความจำของ GPU ได้อย่างเต็มที่ การกำหนดจำนวนวอร์ปต่อบล็อกให้สูงเสมอจึงยังไม่เพียงพอ เนื่องจากการบรรลุแบนด์วิดท์สูงสุดทำได้โดยการใช้ Streaming Multiprocessors (SM) ที่มีอยู่ให้มากที่สุดเท่านั้น ดังนั้น สำหรับจำนวนโทเค็นที่กำหนดและประเภทข้อมูลรวมถึงมิติซ่อนเร้นที่ทราบ เราจึงเลือกมิติตัวเวกเตอร์ต่อเธรดที่เหมาะสมที่สุดก่อน (8 สำหรับ bfloat16, 4 สำหรับ float32) เพื่อให้แน่ใจว่าโค้ดที่สร้างขึ้นใช้คำสั่งโหลด/จัดเก็บขนาด 128 บิต (256 บิตบน Blackwell) เมื่อทราบขนาดเวกเตอร์แล้ว เราสามารถคำนวณจำนวนวอร์ปที่เคอร์เนลต้องการได้ เราพยายามเพิ่มจำนวน SM ให้สูงสุดก่อน: หากมีจำนวนวอร์ปน้อยกว่า SM จะมีการเลือกวอร์ปเดี่ยวต่อบล็อก ซึ่งโดยปกติคือรูปแบบการถอดรหัสของเคอร์เนล หากมีจำนวนวอร์ปเพียงพอ เราจะค่อยๆ เพิ่มจำนวนวอร์ปเป็น 8, 16 หรือ 32 พร้อมกับทำเบนช์มาร์กหาขอบเขตบน ซึ่งหลังจากนั้นเรามักจะไม่เห็นประสิทธิภาพที่ดีขึ้นอีก

แม้ว่า Triton จะทำการ optimization เหล่านี้บางส่วนเบื้องหลัง แต่เราพบว่าเราสามารถบรรลุผลลัพธ์ที่คล้ายคลึงกันได้ด้วยการใช้งาน CuTe layout algebra อย่างระมัดระวัง โดยไม่ต้องถูกจำกัดด้วยข้อบกพร่องของคอมไพล์เลอร์ระดับสูง โดยเฉพาะอย่างยิ่งเมื่อมีการเกี่ยวข้องกับการลดทอนแนวนอนระดับย่อยของวอร์ป (sub-warp horizontal reductions) เราพบว่า CuTeDSL มอบการควบคุมแบบละเอียดที่เราต้องการเพื่อสร้างโค้ดที่เหมาะสมที่สุด

การซิงโครไนซ์กริด (Grid synchronization)

ในบางสถานการณ์ เพื่อกระจายงานไปยังบล็อกต่างๆ ให้มากที่สุดเท่าที่จะเป็นไปได้ สามารถใช้ cooperative_groups::this_grid().sync() เพื่อซิงโครไนซ์บล็อกที่เข้าร่วมทั้งหมดได้ แต่มีข้อจำกัด: เนื่องจากบล็อกไม่สามารถถูกถอดกำหนดการ (de-scheduled) ออกจาก SM ได้เมื่อเริ่มทำงาน เพื่อรับประกันความคืบหน้าไปข้างหน้า จำนวนบล็อกจึงต้องไม่เกินจำนวน SM ที่มีอยู่บนอุปกรณ์ เพื่อให้แน่ใจว่าในทุกจุด บล็อกทั้งหมดสามารถอัปเดตและพอลล์ตำแหน่งหน่วยความจำส่วนกลางที่เกี่ยวข้องกับการซิงโครไนซ์ได้ ข้อจำกัดนี้ถูกบังคับใช้โดย cudaLaunchCooperativeKernel ต้นทุนของ grid barrier อยู่ที่ประมาณ 3us บน Hopper รวมถึงต้นทุนในการฟลัชการดำเนินการหน่วยความจำก่อนหน้าทั้งหมดผ่านหน่วยความจำแบร์ริเออร์โดยปริยายในบางสถานการณ์ โดยรวมแล้ว สิ่งนี้เทียบเท่ากับโอเวอร์เฮดของการแยกเคอร์เนลออกเป็นสองการเรียกใช้งาน ซึ่งนำไปสู่การตัดสินใจออกแบบที่คำนึงถึงประสิทธิภาพ แม้ว่า Programmatic Dependent Launch (PDL) จะสามารถลดโอเวอร์เฮดในการเรียกใช้งานได้ตั้งแต่ Hopper เป็นต้นไป แต่สำหรับเคอร์เนลถอดรหัสแล้ว การหลีกเลี่ยงโอเวอร์เฮดนี้นับว่าคุ้มค่า

หากเวิร์กโหลดตัวใดตัวหนึ่งต้องการบล็อกมากกว่าจำนวน SM ควรจะทำการแยกเคอร์เนล โดยอาศัยแบร์ริเออร์โดยปริยายของการเรียกใช้งานเคอร์เนล ซึ่งโดยหลักแล้วจะเป็นกรณีสำหรับพรีฟิล หากจำนวนบล็อกที่เหมาะสมสำหรับเวิร์กโหลดทั้งหมดพอดีกับกริด จะนิยมใช้ grid barrier เช่นเดียวกับในกรณีของเคอร์เนลถอดรหัสที่ทำงานบนจำนวนโทเค็นที่น้อยกว่า

ใน CuTeDSL การตัดสินใจเหล่านี้สามารถทำได้ในโค้ดการเรียกใช้งานฝั่งโฮสต์ และสามารถคอมโพสและแทรกเคอร์เนลสลับกับ grid barriers ได้ตามความจำเป็น ดังตัวอย่างที่แสดงด้านล่าง ด้วยกลยุทธ์ไฮบริดดังกล่าว เคอร์เนลถอดรหัสสามารถได้รับการปรับ optimization เพื่อให้ได้ Latency ต่ำที่สุด และเคอร์เนลพรีฟิลสามารถปรับ optimization เพื่อให้ได้ throughput สูงสุด แตกต่างจากโซลูชันที่ใช้ CUDA ลอจิกนี้สามารถกำหนดได้อย่างกระชับกว่าและมีโอเวอร์เฮดของเวลาในการคอมไพล์ที่น้อยกว่า

python
@cute.kernel def dispatch_kernel_decode(...): route_tokens(...) grid_barrier(grid_barrier_ptr) copy_tokens_to_experts(...) @cute.kernel def dispatch_kernel_prefill_route(...): route_tokens(...) @cute.kernel def dispatch_kernel_prefill_copy(...): copy_tokens_to_experts(...) @cute.jit def dispatch(...): num_route_blocks = ... num_copy_blocks = ... if num_copy_blocks < NUM_SMS: dispatch_kernel_decode(...).launch( grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1) ) else: dispach_kernel_prefill_route(...).launch( grid=(num_route_blocks, 1, 1) ) dispatch_kernel_prefill_copy(...).launch( grid=(num_copy_blocks, 1, 1) )

ในกรณีของเคอร์เนล MoE all-to-all ของเรา ตามรายละเอียดด้านล่าง การมีความยืดหยุ่นในการสร้างเคอร์เนลที่มีหรือไม่มี grid barriers และการปรับ optimization ที่เวลาคอมไพล์ช่วยให้เราบรรลุการเพิ่มขึ้นของประสิทธิภาพอย่างมีนัยสำคัญผ่าน CuTeDSL

เคอร์เนล CuTeDSL

QK Norm

เคอร์เนลแรกที่เราสร้างขึ้นใน CuTeDSL คือเคอร์เนล RMS norm แบบพิเศษ ซึ่งได้รับการปรับ optimization สำหรับมิติขนาดเล็ก QK norm ซึ่งใช้ในเลเยอร์ self-attention ของโมเดล Qwen รุ่นล่าสุด จะถูกนำไปใช้ในระดับรายเฮด (per-head) บน 64, 128 หรือ 256 อิลิเมนต์ แตกต่างจาก RMS norm ปกติที่ทำงานบนมิติซ่อนเร้นที่เกินกว่าหลายพัน ซึ่งต้องอาศัยการลดทอนแนวนอนผ่านทั้งหน่วยความจำแชร์และรีจิสเตอร์ QK norm จำเป็นต้องมีการลดทอนเฉพาะข้ามวอร์ปเท่านั้น โดยมีความเป็นไปได้ในการบีบเฮดหลายอันให้อยู่ในวอร์ปเดียว สิ่งนี้กลายเป็นกรณีการใช้งานที่เราสามารถทำการปรับ optimization อย่างเข้มข้นที่คอมไพล์ไทม์เพื่อหาการกำหนดโทเค็นต่อวอร์ปที่ดีที่สุดสำหรับมิติเฮดและประเภทข้อมูลทั้งหมด

RMS norm ค่อนข้างเป็นการดำเนินการที่ตรงไปตรงมา แม้ว่าจะต้องมีการหาผลรวมแนวนอนเพื่อคำนวณกำลังสองของนอร์ม (squared norm) ของเวกเตอร์:

cpp
r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

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

จากนั้นเราสามารถสร้างเคอร์เนลแบบพิเศษที่มีบรรทัดโค้ดค่อนข้างน้อย โดยอาศัยการโหลด/จัดเก็บที่ใช้เวกเตอร์โดยอัตโนมัติ และใช้การดำเนินการเฉพาะสถาปัตยกรรมระดับต่ำสำหรับการลดทอนวอร์ป เคอร์เนลจะกำหนดว่าจะต้องทำงานกับเฮดใด (Q หรือ K) และระบุค่าน้ำหนักที่เกี่ยวข้องที่จะนำมาใช้ จากนั้นจะโหลดเวกเตอร์ทั้งหมดจากเฮดและทำการรวมผลรวมย่อยในเครื่องก่อน จากนั้น ตามจำนวนเธรดต่อเฮด สแควร์นอร์มจะถูกคำนวณขึ้น สุดท้าย RMS norm ซึ่งถูกทำเป็นเวกเตอร์โดยอัตโนมัติ จะถูกกำหนดและเขียนกลับไปยังหน่วยความจำ พารามิเตอร์เฉพาะโมเดล เช่น WEIGHT_BIAS และ EPS จะถูกพับค่าคงที่ (constant-folded)

python
@cute.kernel def _qk_norm_kernel( q: cute.Tensor, k: cute.Tensor, q_weight: cute.Tensor, k_weight: cute.Tensor, WEIGHT_BIAS: cutlass.Constexpr[float], EPS: cutlass.Constexpr[float], ) -> None: # Find the block and thread indices. bidx, _, _ = cute.arch.block_idx() bdim, _, _ = cute.arch.block_dim() tidx, _, _ = cute.arch.thread_idx() _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape _, num_k_heads, _ = k.shape head_dim = threads_per_head * vecs_per_thread * vec_size # Determine the thread and the vector to work on. vec_id = (bidx * bdim + tidx) // threads_per_head thread_id = tidx % threads_per_head # Determine the token and head indices. token = vec_id // (num_q_heads + num_k_heads) head_offset = vec_id % (num_q_heads + num_k_heads) is_q = head_offset < num_q_heads # Determine which tensor to work with. thread_coord = (thread_id, None, None) head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord) # Load the head into registers. x_ref = q[head_coord] if is_q else k[head_coord] w = (q_weight if is_q else k_weight)[thread_coord].load() # Compute the squared sum locally. x = x_ref.load() x_f32 = x.to(cute.Float32) x_square = x_f32 * x_f32 x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0) # Sum up across the warp. mask = activemask() start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2) for offset in [start >> i for i in range(5) if (start >> i) > 0]: x_square_sum += cute.arch.shuffle_sync_bfly( x_square_sum, offset=offset, mask=mask, ) # RMS norm. rms = cute.math.rsqrt(x_square_sum / head_dim + EPS) x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type) x_ref.store(x_norm)

เมื่อเทียบกับเคอร์เนล RMS norm ปกติ รูปแบบพิเศษนี้ให้ throughput และ latency ที่ดีขึ้น 2-3 เท่าสำหรับทั้งพรีฟิลและการถอดรหัส

Mixture-of-Experts Dispatch/Combine

MoE dispatch และ combine จะทำการสลับลำดับอินพุตโทเค็นไปยังผู้เชี่ยวชาญ (experts) ที่โทเค็นเหล่านั้นถูกเราท์ไป โดยทำการสะสมแบบถ่วงน้ำหนัก (weighted accumulation) หลังจากที่เคอร์เนล MoE GEMM ทำการคูณเมทริกซ์ที่เกี่ยวข้อง ประสิทธิภาพเป็นสิ่งที่น่าสนใจเป็นพิเศษ เนื่องจากผลิตภัณฑ์ต่างๆ พึ่งพาโมเดล MoE ที่มีขนาดแตกต่างกัน ตามประเภทการปรับใช้งาน เราจึงจำแนกเคอร์เนลต่างๆ เพื่อทำการปรับ optimization ดังนี้:

TP=1: เคอร์เนลโลคอลสำหรับ GPU เดี่ยว Dispatch ทำการสลับโทเค็นไปยังผู้เชี่ยวชาญที่เกี่ยวข้อง และ combine ทำการสะสมค่า ใช้เป็นหลักสำหรับการปรับใช้งาน Qwen3 30B บน H100

DP=1, TP ≠ 1: เคอร์เนลเทนเซอร์พาราเรลแบบโหนดเดี่ยว (singlenode kernel) อินพุตจะถูกทำสำเนาจำลอง, dispatch จะเลือกโทเค็นสำหรับแรงก์ปัจจุบัน และ combine จะทำการสะสมค่าข้ามทุกแรงก์ การสะสมค่าข้าม NVLink หรือ NVLS ใช้สำหรับพรีฟิลเลอร์ขนาดใหญ่

DP ≠ 1, NVLink: เคอร์เนลดาตาพาราเรลพร้อมเทนเซอร์พาราเรลที่เป็นทางเลือก แต่ละแรงก์ส่งโทเค็นไปยังแรงก์อื่นๆ ทั้งหมดที่โฮสต์ผู้เชี่ยวชาญเฉพาะ ข้อมูลการเราท์จะถูกรวมเข้าด้วยกัน และโทเค็นจะถูกถ่ายโอนผ่าน NVIDIA NVLink ใช้สำหรับการปรับใช้งานแบบดาตาพาราเรลบน H200 หรือ GB200

DP ≠ 1, InfiniBand: เคอร์เนลดาตาพาราเรลและเทนเซอร์พาราเรลแบบ Peer-to-Peer ที่ทำงานข้ามทั้ง NVLink และ NVIDIA Quantum InfiniBand ใช้สำหรับการปรับใช้งาน Hopper แบบหลายโหนด

โลคอล (Local)

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

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

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

โหนดเดี่ยว (Singlenode)

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

เนื่องจากจำนวนแรงก์มักจะน้อยกว่าจำนวนผู้เชี่ยวชาญที่โทเค็นถูกเราท์ไป ในกรณีโดยทั่วไปจึงสามารถสันนิษฐานได้ว่าโทเค็นแต่ละตัวมีสำเนาอย่างน้อยหนึ่งชุดถูกเราท์ไปยังแต่ละแรงก์ ในสถานการณ์นี้ เคอร์เนล combine โลคอลสามารถนำกลับมาใช้ใหม่เพื่อรวมสำเนาโลคอลของโทเค็นเข้าด้วยกัน ตามด้วยการทำ all-reduce เพื่อรวมค่าผลรวมข้ามแรงก์ที่เข้าร่วม ในกรณีของเคอร์เนลพรีฟิล เราจะเรียกใช้เคอร์เนลแยกต่างหาก อย่างไรก็ตาม การถอดรหัสจะผสานการดำเนินการทั้งสองเข้าเป็นการเรียกใช้งานครั้งเดียว

สำหรับการทำ all-reduce เราใช้ NVLink แบบปกติหากมีแรงก์เข้าร่วมเพียงสองแรงก์ และจะสลับไปใช้ NVLS เมื่อมีแรงก์ตั้งแต่ 4 แรงก์ขึ้นไป ซึ่งจะช่วยเพิ่ม throughput ได้ 10-20% ต้องขอบคุณ DSL ทำให้เราสามารถแชร์ลอจิกการลดทอนส่วนใหญ่และการแลกเปลี่ยนบัฟเฟอร์ peer ระหว่าง MoE และเคอร์เนล all-reduce ปกติได้ นี่คือสถานการณ์ที่การควบคุมฮาร์ดแวร์อย่างใกล้ชิดมีประโยชน์ เนื่องจากจำเป็นต้องมีการดำเนินการ DSL ผ่าน multimem.ld_reduce และ multimem.st สำหรับ NVLS การใช้งานการลดทอนเองนั้นดำเนินตามตัวอย่าง all_reduce จากการทดสอบ NCCL เป็นส่วนใหญ่

ดาตาพาราเรล (Data Parallel)

เคอร์เนล Data Parallel อาศัย NVLink เพื่อรวมโทเค็นสำหรับอุปกรณ์หลายเครื่อง ในเฟส dispatch ก่อนที่จะแลกเปลี่ยนโทเค็น เพียร์จะทำเราท์ในเครื่องก่อนเพื่อระบุออฟเซ็ตเป้าหมายบนแต่ละแรงก์สำหรับโทเค็นแต่ละตัวที่กำลังส่ง จากนั้นจึงเริ่มคัดลอกโทเค็นไปยังบัฟเฟอร์รับที่จัดสรรไว้สำหรับผู้ส่งแต่ละรายเพื่อหลีกเลี่ยงข้อขัดแย้ง การเขียนแบบ release-store ต่อโทเค็นจะถูกใช้เพื่อส่งสัญญาณว่าเสร็จสิ้น ที่ฝั่งผู้รับ ข้อมูลการเราท์จะถูกรวมเข้าด้วยกันเพื่อให้สามารถย่อโทเค็นทั้งหมดให้อยู่ในบัฟเฟอร์เดียวสำหรับ MoE GEMM ผู้รับจะทำการพอลล์แฟล็กที่เขียนโดย peer ด้วย acquire-load เพื่อรอให้การเขียนเสร็จสมบูรณ์

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

เราได้สร้างไลบรารีของ PTX wrappers เพื่อเปิดเผย release-acquire loads และ stores ที่ขอบเขตระดับระบบและระดับอุปกรณ์ เพื่อใช้งานเคอร์เนลนี้ ด้วยการปรับปรุงความละเอียดของแฟล็กการซิงโครไนซ์และการจัดสรรบล็อกเพิ่มเติมเพื่อส่งและรับโทเค็นผ่าน NVLink เราจึงบรรลุประสิทธิภาพที่ดีขึ้น 2 เท่าเมื่อเทียบกับเคอร์เนล CUDA ก่อนหน้าของเรา

Peer-to-Peer

ก่อนหน้านี้ เคอร์เนล Peer-to-Peer ของเราทำประสิทธิภาพระดับล้ำสมัยบน GPU Hopper ที่จับคู่กับอแดปเตอร์ ConnectX-7 InfiniBand ตั้งแต่นั้นมาเราได้พอร์ตเคอร์เนล combine ไปยัง CuTeDSL เพื่อให้ได้ประสิทธิภาพที่เพิ่มขึ้นอีก เนื่องจากเวอร์ชัน CUDA อาศัยการปรับแต่งแบบคอมไพล์ไทม์อย่างมาก จึงพิสูจน์ได้ว่าเป็นคอขวดในเวลาในการคอมไพล์ภายในคลังโค้ดของเรา ซึ่งบีบให้เราต้องละเว้นพารามิเตอร์บางตัว การย้ายไปสู่ DSL ที่คอมไพล์แบบ JIT ทำให้เราสามารถปรับแต่งสำหรับพารามิเตอร์ที่มากขึ้นและปรับแต่งการจัดสรรกริด ทำให้เกิดความแตกต่างเพิ่มเติมระหว่างพรีฟิลและการถอดรหัส ซึ่งช่วยเพิ่ม Latency ได้ประมาณ 3%

การสนับสนุน Blackwell

จากมุมมองของเคอร์เนล ROSE การสนับสนุน Blackwell แบ่งออกเป็นสองประเภท: เคอร์เนลที่ไม่ขึ้นกับสถาปัตยกรรมและเคอร์เนลที่ขึ้นกับสถาปัตยกรรม เคอร์เนลที่ขึ้นกับสถาปัตยกรรมซึ่งอาศัย Tensor cores จำเป็นต้องมีการทำงานที่ชัดเจนเพื่อปรับตัว อย่างไรก็ตาม ด้วยการใช้ DSL ทำให้เคอร์เนลที่ไม่ขึ้นกับสถาปัตยกรรมสามารถปรับตัวได้อย่างง่ายดาย

GPU Blackwell มี SM สูงสุดถึง 160 หน่วย ตรงกันข้ามกับ 132 หน่วยของ Hopper สำหรับเคอร์เนลส่วนใหญ่ สิ่งนี้คือค่าคงที่ของคอมไพล์ไทม์ และลอจิกในโค้ดตัวเปิดใช้ฝั่งโฮสต์จะนำปัจจัยนี้มาพิจารณาเพื่อกำหนดเคอร์เนลที่เหมาะสมที่สุด (พรีฟิลเทียบกับการถอดรหัส) ที่จะรัน นอกจากนี้ Blackwell ยังแนะนำคำสั่งโหลดขนาด 256 บิต ซึ่งสามารถปล่อยออกมาได้อย่างง่ายดายโดยการเพิ่มความกว้างของเวกเตอร์ (มิติล่าสุดของเทนเซอร์ CuTe) เพื่อปรับปรุงการใช้แบนด์วิดท์หน่วยความจำให้ดีขึ้น

Blackwell แนะนำการสนับสนุนฮาร์ดแวร์สำหรับรูปแบบการควอนไทเซชันจำนวนมาก (MXFP8, MXFP4, NVFP4) การสนับสนุนสิ่งเหล่านี้ถูกนำมาใช้เป็นหลักโดยการปรับแต่งเคอร์เนลที่มีอยู่สำหรับการทำสเกลแบบบล็อกต่อบล็อกของ FP8 การสนับสนุน JIT มีความสำคัญอย่างยิ่งในที่นี้ เนื่องจากโหมดใหม่สามารถได้รับการสนับสนุนได้โดยไม่ทำให้สเปซการกำหนดค่าขณะคอมไพล์ระเบิดขึ้น ในอดีต เราพบว่า Triton มีความอึดอัดในการแสดงการลดทอนแนวนอนข้ามอิลิเมนต์จำนวนน้อย เกินครึ่งหรือหนึ่งในสี่ของวอร์ป การควบคุมที่ละเอียดยิ่งขึ้นเกี่ยวกับการลดทอนเป็นสิ่งสำคัญยิ่งเพื่อให้บรรลุประสิทธิภาพสูงสุดเมื่อทำการคำนวณค่าสูงสุดสัมบูรณ์ข้าม 16 หรือ 32 อิลิเมนต์สำหรับ NVFP4 และ MXFP8/MXFP4 ตามลำดับ

การสร้างเคอร์เนลด้วย AI (AI-Assisted Kernel Generation)

ในฐานะที่เป็นภาษาค่อนข้างใหม่ โมเดล LLM จึงแทบไม่มีการสัมผัสกับโค้ด CuTeDSL เลยในช่วงการฝึกอบรมล่วงหน้า นอกจากนี้ สไตล์ของโค้ดปัจจุบันของเรา ซึ่งมีเคอร์เนลที่ได้รับการปรับแต่งสำหรับแอคติเวชันและ norm นั้นแตกต่างเล็กน้อยจากตัวอย่าง CUTLASS หรือ Quack ที่มีอยู่ ซึ่งได้รับการปรับแต่งอย่างมากสำหรับคำสั่ง TMA และการคูณเมทริกซ์เฉพาะสถาปัตยกรรม การไปถึงจุดที่ LLM สามารถช่วยเหลือในการเขียนเคอร์เนลได้นั้นเป็นกระบวนการทีละขั้นตอน

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

บทสรุปและงานในอนาคต

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

ณ 1 ในขณะที่เขียนนี้ ช่องว่างด้านประสิทธิภาพระหว่าง CUTLASS C++ และเคอร์เนล CuTeDSL นั้นมีน้อยมาก ซึ่งเปิดโอกาสที่น่าตื่นเต้นในการคอมโพสและผสานเคอร์เนล GEMM ประสิทธิภาพสูงเข้ากับการสื่อสารระหว่างอุปกรณ์, นอร์มัลไลเซชัน หรือแอคติเวชัน ด้วยการมุ่งมั่นทุ่มเทให้กับ DSL นี้ เราจะสามารถดึงประสิทธิภาพเพิ่มเติมจากอุปกรณ์ของเราเพื่อปรับปรุง throughput ข้ามเวิร์กโหลดของเรา ตั้งแต่เอ็มเบดดิ้งไปจนถึง LLM

เอกสารอ้างอิง

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

การอ้างอิง

BibTeX
@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}