การปรับแต่งการอนุมานบนอุปกรณ์ให้เหมาะสมสำหรับ Apple Silicon

การประมวลผลแบบไฮบริดบน Apple silicon ประสานงานการทำงานระหว่างความฉลาดระดับฟรอนเทียร์ในคลาวด์และโมเดลภายในเครื่องบน Mac โมเดลบนคลาวด์จัดการงานวิจัยและการใช้เหตุผล ในขณะที่โมเดลภายในเครื่องทำงานร่วมกับไฟล์ส่วนตัวและแอปบน Mac

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

Lily ซึ่งเป็นเอ็นจิ้นการอนุมานภายในเครื่องที่มีน้ำหนักเบาของเรา ถูกสร้างขึ้นมาโดยเฉพาะสำหรับ Apple silicon และ Qwen3.6-35B-A3B โดยมีการปรับแต่งแยกต่างหากสำหรับ prefill และ decode เอ็นจิ้นนี้จะเปิดซอร์สโค้ดในเร็วๆ นี้

บทนำ

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

Qwen3.6-35B-A3B เป็นโมเดลแบบกระจัดกระจายและไฮบริด: โดยใช้การกำหนดเส้นทางแบบ mixture-of-experts (MoE) และผสมผสานสถานะการเรียกซ้ำขนาดคงที่เข้ากับ full attention ทางเลือกด้านสถาปัตยกรรมเหล่านี้ช่วยลดปริมาณการคำนวณที่ต้องใช้ แต่ก็สร้างภาระงานที่ไม่สม่ำเสมอเช่นกัน โทเค็นจะถูกส่งไปยังน้ำหนักของผู้เชี่ยวชาญที่แตกต่างกัน และสถานะการเรียกซ้ำจะมีลักษณะเป็นลำดับ

MLX-LM เลือกเคอร์เนลที่ปรับแต่งแล้วสำหรับเฟสการอนุมานและรูปร่างภาระงานทั่วไปอยู่แล้ว แต่การทำงานที่นำกลับมาใช้ใหม่ได้ของมันจะต้องรองรับสถาปัตยกรรมโมเดลจำนวนมาก เอ็นจิ้นที่อุทิศให้กับ Qwen สามารถเชี่ยวชาญเฉพาะทางในระดับโมเดลและรันไทม์ โดยประสานเคอร์เนล การเคลื่อนย้ายข้อมูล และการจัดตารางเวลารอบโครงสร้างคงที่ของโมเดล

Lily นำการปรับแต่งเฉพาะทางนี้ไปใช้แบบ end-to-end ในกระบวนการเดียว รันไทม์ Rust โหลดเช็คพอยต์โมเดลและจัดการสถานะเซสชันและลูปการสร้าง, API การสนทนาแชทที่เข้ากันได้กับ OpenAI ยอมรับคำขอและสตรีมโทเค็น, และเคอร์เนล Metal แบบกำหนดเองทำหน้าที่รันการทำงานเฉพาะของ Qwen ไม่มีทั้ง PyTorch หรือ MLX อยู่ในเส้นทางการทำงาน

เราวัดประสิทธิภาพของ prefill และ decode แยกกัน อัตราการส่งข้อมูล prefill จะจับความเร็วที่เอ็นจิ้นประมวลผลพรอมต์ ส่วนอัตราการส่งข้อมูล decode จะจับความเร็วที่เอ็นจิ้นสร้างโทเค็นผลลัพธ์

เราทำการทดสอบประสิทธิภาพ Qwen3.6-35B-A3B บน MacBook Pro เครื่องเดียวที่ขับเคลื่อนด้วย M5 Max ซึ่งมี 40-core GPU และหน่วยความจำแบบรวม (unified memory) 128 GB ผ่านความยาวพรอมต์ 10 ระดับสำหรับ prefill และความยาวบริบท 10 ระดับสำหรับ decode ตั้งแต่ 256 ถึง 128K โทเค็น (K = 1,024) เอ็นจิ้นนี้มีอัตราการส่งข้อมูล prefill เฉลี่ย 1.23 เท่า และอัตราการส่งข้อมูล decode 1.35 เท่าของ MLX-LM ที่พรอมต์ 4K โทเค็นและบริบท decode 4K โทเค็น เอ็นจิ้นแบบกำหนดเองนี้ทำความเร็ว prefill ได้ 5,749.9 โทเค็นต่อวินาที และ decode ได้ 186.6 โทเค็นต่อวินาที เมื่อเทียบกับ 4,737.5 และ 140.9 ของ MLX-LM ในเซสชันแบบหลายรอบ (multi-turn) การประหยัดเวลาเหล่านี้จะสะสมเพิ่มขึ้นในทุกๆ การเรียกใช้งานโมเดล

ถัดไป เราจะอธิบายว่าสถาปัตยกรรมของ Qwen สร้างโอกาสในการปรับแต่งเฉพาะโมเดลบน Apple silicon ได้อย่างไร จากนั้นเราจะอธิบายการเปลี่ยนแปลงของ prefill และ decode ที่ตามมา นอกจากนี้ เรายังครอบคลุมถึงจุดที่การปรับแต่งเพิ่มเติมหยุดให้ผลคุ้มค่า ก่อนจะปิดท้ายด้วยการเปรียบเทียบแบบ end-to-end กับ MLX-LM

โอกาสในการปรับแต่งเฉพาะสำหรับ Qwen บน Apple silicon

Qwen สร้างรูปร่างภาระงานที่แตกต่างกัน 3 รูปแบบ

Qwen3.6-35B-A3B มีพารามิเตอร์ 3.5 หมื่นล้านตัวแต่เปิดใช้งานเพียงประมาณ 3 พันล้านตัวสำหรับโทเค็นแต่ละตัว ตัวจัดเส้นทางจะให้คะแนนเครือข่ายย่อยของผู้เชี่ยวชาญ 256 เครือข่ายและเลือก 8 เครือข่าย ควบคู่ไปกับผู้เชี่ยวชาญที่แชร์หนึ่งรายซึ่งประมวลผลทุกๆ โทเค็น การออกแบบ MoE แบบกระจัดกระจายนี้ช่วยลดการคำนวณแต่สร้างงานที่ไม่สม่ำเสมอ: ผู้เชี่ยวชาญได้รับจำนวนโทเค็นที่แตกต่างกัน และโทเค็นแต่ละตัวต้องการน้ำหนักจากชุดค่าผสมของผู้เชี่ยวชาญที่แตกต่างกัน

Qwen ยังผสมผสานเลเยอร์ full-attention 10 เลเยอร์เข้ากับเลเยอร์ Gated DeltaNet 30 เลเยอร์ เลเยอร์ทั้งสองประเภทนี้เก็บรักษาข้อมูลก่อนหน้าไว้ในลักษณะที่แตกต่างกัน

เลเยอร์ attention ใช้ grouped-query attention (GQA) Qwen มีหัว query 16 หัวและหัว key-value (KV) 2 หัว โดยหัว query แปดหัวจะแชร์หัว KV ร่วมกัน การแชร์นี้ทำให้แคช KV มีขนาดเล็กลงและช่วยให้นำข้อมูลที่แคชไว้กลับมาใช้ซ้ำข้ามหัว query ต่างๆ ได้ แคชยังคงจัดเก็บ key และ value ใหม่สำหรับทุกโทเค็น ดังนั้นแต่ละขั้นตอนการ decode จึงอ่านข้อมูลเพิ่มขึ้นเมื่อบริบทเติบโตขึ้น

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

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

Apple silicon มีเส้นทางที่แตกต่างกันสำหรับภาระงานที่แตกต่างกัน

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

Apple silicon วาง CPU และ GPU ไว้เบื้องหลังหน่วยความจำแบบรวม (unified memory) ซึ่งเป็นพูลหน่วยความจำทางกายภาพเดียวที่ทั้งสองสามารถเข้าถึงได้ วิธีนี้ช่วยให้โมเดลยังคงอยู่ในหน่วยความจำโดยไม่ต้องรักษาส่วนสำเนาแยกต่างหากบน GPU แต่ไม่ได้ทำให้การเคลื่อนย้ายข้อมูลเป็น 0 การอ่านน้ำหนักและค่าระหว่างกลางยังคงบริโภคแบนด์วิดท์หน่วยความจำ ในขณะที่รีจิสเตอร์และพื้นที่จัดเก็บข้อมูลบนชิปอื่นๆ จะเร็วกว่าแต่มีขนาดเล็กกว่ามาก

M5 GPU ยังมีเส้นทางการประมวลผลที่แตกต่างกัน เลเยอร์เชิงเส้นของ prefill ใช้การคูณเมทริกซ์-เมทริกซ์ทั่วไป (GEMM) โดยการใช้เมทริกซ์น้ำหนักกับหลายแถวพร้อมกัน GEMM ที่เข้ากันได้สามารถใช้ Neural Accelerator ในคอร์ GPU แต่ละคอร์ผ่าน การทำงานของเทนเซอร์ Metal 4 ในทางกลับกัน Decode แบบ batch-1 จะใช้การคูณเมทริกซ์-เวกเตอร์ทั่วไป (GEMV) โดยการใช้น้ำหนักเดียวกันกับหนึ่งแถว เนื่องมีการใช้น้ำหนักซ้ำน้อย GEMV จึงถูกจำกัดหลักๆ ด้วยแบนด์วิดท์หน่วยความจำ และเหมาะกับเวกเตอร์ ALU ของ GPU มากกว่า Neural Accelerator ที่ออกแบบมาสำหรับการทำงานของเมทริกซ์ที่มีการนำข้อมูลกลับมาใช้ซ้ำสูงกว่า

เส้นทางการทำงานเหล่านี้ไม่ได้มีแค่ใน Lily เท่านั้น MLX ทำงานบนหน่วยความจำแบบรวมเดียวกันและเลือกเคอร์เนลเมทริกซ์และเวกเตอร์ที่ปรับแต่งแล้วตามรูปร่างภาระงาน การใช้งาน Qwen ของ MLX-LM ทำการจัดกลุ่มงานผู้เชี่ยวชาญ, ประเมิน Gated DeltaNet ด้วยเคอร์เนล Metal แบบเรียกซ้ำที่ผสานรวมกัน และใช้ attention ที่ตระหนักถึง GQA อยู่แล้ว ความสามารถเหล่านี้คือจุดเริ่มต้นร่วมกันสำหรับการอนุมาน Qwen ที่มีประสิทธิภาพบน Apple silicon

กลยุทธ์การปรับแต่ง

ขอบเขตที่แคบลงของ Lily ช่วยให้สามารถประสานเส้นทางการทำงานที่ใช้ร่วมกันเหล่านี้รอบๆ สถาปัตยกรรมและมิติข้อมูลที่แน่นอนของ Qwen ได้ โดยใช้เส้นทาง GPU เฉพาะเฟส, จับคู่ภาระงานของผู้เชี่ยวชาญ, การเรียกซ้ำ และ attention ของ Qwen เพื่อลดการเคลื่อนย้ายข้อมูลให้เหลือน้อยที่สุด และเลือกเคอร์เนลและเลย์เอาต์จากรูปร่างภาระงานที่วัดได้ กลยุทธ์นี้มีสามส่วนด้วยกัน:

1. จับคู่เส้นทาง GPU ให้เข้ากับเฟสการอนุมาน ใช้การทำงานที่เน้นเมทริกซ์เมื่อ prefill สามารถนำน้ำหนักกลับมาใช้ซ้ำข้ามแถวจำนวนมากได้ และใช้การทำงานที่เน้นเวกเตอร์เมื่อการ decode แบบ batch-1 ประมวลผลทีละแถว
2. แมปโครงสร้างของ Qwen ลงบน GPU พร้อมลดการเคลื่อนย้ายข้อมูลให้น้อยที่สุด เก็บรักษาน้ำหนักให้อยู่ในรูปแบบบีบอัดจนกว่าจะใช้งาน, จัดระเบียบงานผู้เชี่ยวชาญที่มีการกำหนดเส้นทางโดยไม่ต้องกลับไปที่ CPU, คงสถานะ Gated DeltaNet ไว้บนชิปตลอดการสแกนแบบเรียกซ้ำ, และนำข้อมูล KV ที่แชร์โดย grouped-query attention กลับมาใช้ใหม่
3. ปรับแต่งเคอร์เนลให้เข้ากับรูปร่างภาระงาน ภายในแต่ละเฟส ให้เลือกขนาดไทล์, เลย์เอาต์การทำงาน และเส้นทาง attention จากจำนวนแถวที่มีอยู่, การกระจายตัวของแถวไปยังผู้เชี่ยวชาญ, มิติข้อมูลของการทำงาน และความยาวบริบทปัจจุบัน

ส่วนต่อไปนี้จะอธิบายทางเลือกเหล่านี้ สำหรับการปรับแต่งที่ได้รับการประเมินในการทดสอบ ablation ที่ตรงกันบน M5 Max เราประเมินผลกระทบโดยการเปรียบเทียบการกำหนดค่าเอ็นจิ้นที่เหมือนกันทุกประการ ยกเว้นส่วนของการปรับแต่งที่อยู่ภายใต้การศึกษา เนื่องจาก การทดลองเหล่านี้เป็นการเปรียบเทียบเวอร์ชันต่างๆ ของเอ็นจิ้นของเรากับตัวมันเอง จึงเป็นการอธิบายกลไกมากกว่าการแยกแยะผลลัพธ์สุดท้ายเทียบกับ MLX-LM

Prefill: นำน้ำหนักกลับมาใช้ซ้ำและให้การกำหนดเส้นทางอยู่บน GPU

Prefill เปิดเผยแถวโทเค็นจำนวนมากพร้อมกัน แต่ Qwen จัดเส้นทางแถวเหล่านั้นไปยังผู้เชี่ยวชาญอย่างไม่สม่ำเสมอและอัปเดตสถานะการเรียกซ้ำตลอดลำดับ การปรับแต่งของมันแบ่งออกเป็นสามกลุ่ม: จัดระเบียบงานผู้เชี่ยวชาญแบบกระจัดกระจายรอบๆ แถวที่มีการกำหนดเส้นทาง, รักษาการสแกน Gated DeltaNet ไว้บนชิป, และแบ่งพรอมต์ขนาดยาวออกเป็นส่วนย่อยที่มีขอบเขตจำกัด

ปรับปรุงการคำนวณของผู้เชี่ยวชาญแบบกระจัดกระจาย (sparse expert computation)

แปลงน้ำหนักกลับเป็นค่าเดิม (dequantize) ระหว่างการคูณเมทริกซ์

เช็คพอยต์ Qwen3.6-35B-A3B ใช้การควอนตาร์เซชันแบบแอฟฟีนกลุ่มละ (groupwise affine 4-bit quantization) น้ำหนักแต่ละตัวจะถูกจัดเก็บเป็นรหัสจำนวนเต็ม 4 บิต ในขณะที่น้ำหนักทุกๆ กลุ่ม 64 ตัวจะใช้สเกลและไบอัส bfloat16 ร่วมกันเพื่อสร้างค่าของมันขึ้นมาใหม่ วิธีนี้ช่วยลดโมเดลที่มีพารามิเตอร์ 3.5 หมื่นล้านตัวจากน้ำหนัก bfloat16 ขนาดประมาณ 70 GB ให้เหลือเช็คพอยต์ขนาด 19.4 GB ทำให้สามารถเก็บโมเดลไว้ใน Mac ได้จริง

การทำงานของเทนเซอร์ Metal 4 ที่ใช้สำหรับการคูณเมทริกซ์จะใช้โอเปอแรนด์ bfloat16 แทนที่จะเป็นตัวแทน 4 บิตที่แพ็คไว้ ก่อนการคูณ GPU จะต้องสร้างน้ำหนักในรูปแบบ bfloat16 ขึ้นมาใหม่ Grouped GEMM ที่ปรับแต่งแล้วใน Lily จะทำการแปลงนี้ทีละไทล์น้ำหนักขนาดเล็ก และเก็บผลลัพธ์ไว้ในหน่วยความจำ threadgroup บนชิปนานพอที่จะคูณมันเข้ากับแถวการเปิดใช้งานที่มีการกำหนดเส้นทาง การสะสมจะใช้จุดลอยตัว 32 บิต และเอาต์พุตจะถูกเขียนในรูปแบบ bfloat16 อาร์เรย์น้ำหนักที่ขยายออกจนสมบูรณ์จะไม่ถูกสร้างขึ้นในหน่วยความจำแบบรวมเลย

ในการทำ ablation การแปลงค่ากลับ (dequantization) จะทำงานเป็นกระบวนการแยกต่างหาก: มันจะขยายน้ำหนัก 4 บิตให้เป็นอาร์เรย์ bfloat16 ในหน่วยความจำแบบรวม หลังจากนั้นเคอร์เนลเมทริกซ์จะอ่านอาร์เรย์นั้นกลับเข้ามา ที่พรอมต์ 512 โทเค็น การย้ายการแปลงค่ากลับเข้าไปใน grouped GEMM ช่วยเพิ่มอัตราการส่งข้อมูล prefill แบบ end-to-end ขึ้น 77.4% โดยการกำจัดการเขียนและการอ่านระหว่างกลางนี้

ให้การกำหนดเส้นทางผู้เชี่ยวชาญอยู่บน GPU

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

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

ที่พรอมต์ 512 โทเค็น การเปิดใช้งานการกำหนดเส้นทางที่อยู่บน GPU ช่วยเพิ่ม prefill แบบ end-to-end ขึ้น 89% สิ่งนี้ยังแสดงให้เห็นว่าเหตุผลที่จำนวนเคอร์เนลอย่างเดียวอาจทำให้เข้าใจผิดได้: เส้นทางที่เร็วกว่านั้นจะเรียกใช้เคอร์เนลมากกว่าแต่ไม่ต้องรอ CPU ภายในเลเยอร์เลย

จับคู่ขนาดไทล์กับภาระงานของผู้เชี่ยวชาญ

ที่พรอมต์ 2K โทเค็น การส่งโทเค็นแต่ละรายการไปยังผู้เชี่ยวชาญ 8 จาก 256 รายการจะสร้างการกำหนดค่าโทเค็นและผู้เชี่ยวชาญ 16,384 รายการ หรือเฉลี่ย 64 แถวการเปิดใช้งานต่อผู้เชี่ยวชาญ การกระจายตัวจริงมีความไม่สม่ำเสมอ: ผู้เชี่ยวชาญบางรายได้รับหลายแถว ในขณะที่บางรายได้รับน้อย

Grouped GEMM จะแบ่งเอาต์พุตของผู้เชี่ยวชาญแต่ละรายออกเป็นไทล์ ซึ่งเป็นบล็อกสี่เหลี่ยมผืนผ้าขนาดเล็กของผลลัพธ์การคูณเมทริกซ์ ไทล์แต่ละไทล์จะถูกกำหนดให้กับ threadgroup ของ GPU หนึ่งกลุ่ม บน GPU ของ Apple silicon ตัว threadgroup จะประกอบด้วย simdgroup หนึ่งกลุ่มขึ้นไป โดยแต่ละกลุ่มประกอบด้วยเธรด 32 เธรดที่ดำเนินการคำสั่งพร้อมกันในจังหวะเดียวกัน (lockstep)

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

การทำ ablation กำหนดขนาดไทล์ไว้ที่ 16 แถว เมื่อเทียบกับตัวควบคุมนั้น การเปิดใช้งานไทล์ขนาด 32 แถวที่มี simdgroup สี่กลุ่มช่วยปรับปรุง prefill แบบ end-to-end ขึ้น 13.2% ที่พรอมต์ 2K โทเค็น

เก็บสถานะการเรียกซ้ำไว้บนชิป

ในระหว่าง prefill เลเยอร์ Gated DeltaNet แต่ละเลเยอร์จะสแกนพรอมต์ตามลำดับพร้อมกับส่งสถานะการเรียกซ้ำไปข้างหน้า เมื่อปิดใช้งานการอยู่ในรีจิสเตอร์ การทำ ablation จะใช้การสแกนแบบ blockwise ที่พรอมต์ 2K โทเค็น เส้นทางนั้นจะเคลื่อนย้ายสถานะ 256 MiB (mebibytes) ต่อเลเยอร์ และหยุดการทำงานของเธรดที่ร่วมมือกันซ้ำแล้วซ้ำเล่าที่จุดอุปสรรค (barriers) ซึ่งเป็นจุดซิงโครไนซ์ที่เธรดที่เข้าร่วมทั้งหมดต้องรอซึ่งกันและกัน

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

สถานะและเกตของมันใช้รูปแบบจุดลอยตัว 32 บิต เนื่องจากข้อผิดพลาดในการปัดเศษเล็กน้อยจะสะสมทบกันผ่านการอัปเดตแบบลำดับ การกระตุ้นเตือนของ query และ key ยังคงอยู่ในรูปแบบ bfloat16

ที่พรอมต์ 2K โทเค็น การเปิดใช้งานการสแกนในรีจิสเตอร์ช่วยปรับปรุง prefill แบบ end-to-end ขึ้น 5.6% Expert GEMM มีสัดส่วนประมาณ 90% ของเวลา prefill การสแกนแบบลำดับไม่ได้เปิดเผยงานเมทริกซ์ที่สามารถนำกลับมาใช้ใหม่ได้มากพอที่จะได้รับประโยชน์จาก Neural Accelerators

จำกัดหน่วยความจำชั่วคราวด้วยการแบ่งพรอมต์เป็นส่วนย่อย (prompt chunking)

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

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

Prefill แบบแบ่งส่วนย่อยเป็นที่นิยมในหลายเอ็นจิ้นและมีความสำคัญอย่างยิ่งสำหรับการให้บริการการสนทนาแบบหลายรอบ (multi-turn trajectories) ขนาดยาวในสภาพแวดล้อมที่มีหน่วยความจำจำกัดเหล่านี้ เวลา prefill ทั้งหมดสำหรับเลเยอร์ attention ยังคงเพิ่มขึ้นตามกำลังสองของความยาวพรอมต์ พร้อมกับมีโอเวอร์เฮดเพิ่มเติมบางส่วนจากการโหลด KV ซ้ำของส่วนย่อยก่อนหน้า

Decode: ลดจำนวนไบต์ที่เคลื่อนย้ายต่อโทเค็น

Decode แบบ Batch-1 จะประมวลผลแถวใหม่ทีละแถว เนื่องจากมีการใช้น้ำหนักซ้ำน้อย อัตราการส่งข้อมูลจึงขึ้นอยู่กับจำนวนไบต์ที่เอ็นจิ้นเคลื่อนย้ายสำหรับโทเค็นแต่ละตัวเป็นหลัก การเปลี่ยนแปลงในการ decode แบ่งออกเป็นสี่กลุ่ม: ปรับแต่งเส้นทางน้ำหนักแถวเดียว, รักษาแต่ละขั้นตอนให้อยู่บน GPU, ลดการรับส่งข้อมูลระหว่างกลางและสถานะ, และอ่านแคช attention อย่างมีประสิทธิภาพ

ปรับแต่งเส้นทางน้ำหนักแถวเดียว

MLX ทำการส่งงานแถวเดียวไปยังเคอร์เนลเมทริกซ์-เวกเตอร์เฉพาะทางอยู่แล้ว เนื่องจาก Lily ไม่ได้ใช้ MLX รันไทม์แบบกำหนดเองจึงต้องจัดเตรียมกลยุทธ์พื้นฐานที่เหมือนกัน GEMV แบบขนานแถวของเราออกแบบมาสำหรับแถวการเปิดใช้งานแถวเดียว Simdgroup จะร่วมมือกันประมวลผลผลลัพธ์พร้อมกับอ่านส่วนต่างๆ ของเมทริกซ์น้ำหนักแบบขนานกันไปด้วย

ให้แต่ละขั้นตอนการ decode อยู่บน GPU

ให้การส่งต่อโทเค็นอยู่บน GPU

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

ให้งาน GPU อิสระทำงานซ้อนทับกัน

ในขั้นตอนการ decode แบบ batch-1 ที่บันทึกไว้ครั้งหนึ่ง การสร้างโทเค็นได้เรียกใช้เคอร์เนล GPU 795 ตัว การพึ่งพากันของพวกมันสร้างเป็นขั้นตอนแบบลำดับ 555 ขั้นตอน ทำให้เคอร์เนลบางตัวสามารถทำงานพร้อมกันได้อย่างอิสระ อย่างไรก็ตาม โหมดการทำงานแบบอนุกรมของ Metal ทำให้ทุกเคอร์เนลทำงานตามลำดับ

เส้นทาง decode ที่ปรับแต่งแล้วจะบันทึกการพึ่งพาข้อมูลจริงใน Metal pass แบบพร้อมกัน การเรียกใช้เคอร์เนลอิสระสามารถทำงานพร้อมกันได้เมื่อทรัพยากร GPU เอื้ออำนวย อุปสรรค (barrier) จะถูกแทรกเข้ามาเฉพาะเมื่อการทำงานภายหลังต้องการผลลัพธ์จากก่อนหน้าเท่านั้น

ลดการรับส่งข้อมูลระหว่างกลางและสถานะ

เคอร์เนลแยกมักจะสร้างค่าระหว่างกลางขึ้นมา: เคอร์เนลหนึ่งเขียนผลลัพธ์ชั่วคราวลงในหน่วยความจำ และเคอร์เนลถัดไปจะอ่านผลลัพธ์นั้นกลับขึ้นมา เส้นทาง decode ที่ปรับแต่งแล้วจะผสานโซ่การทำงานสี่ชุดเข้าด้วยกัน: การฉายภาพอินพุตของผู้เชี่ยวชาญสองชุดพร้อมกับการกระตุ้นเตือนแบบเกต (gated activation), การฉายภาพเอาต์พุตของผู้เชี่ยวชาญพร้อมกับคะแนนการกำหนดเส้นทางและผลลัพธ์ของผู้เชี่ยวชาญที่แชร์, การเตรียม query และ key ก่อนทำ attention, และการอัปเดตแบบเรียกซ้ำพร้อมกับการปรับมาตรฐาน (normalization) เคอร์เนลที่ผสานรวมแต่ละตัวจะเก็บค่าชั่วคราวไว้ในรีจิสเตอร์แทนที่จะส่งผ่านหน่วยความจำ

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

อ่านแคช attention อย่างมีประสิทธิภาพ

รวมการอ่านแคช attention เข้าด้วยกัน (coalesce)

Attention จะอ่าน key และ value จากแคช KV ในทุกขั้นตอนการ decode ในการทำ ablation เธรด GPU ที่อยู่ข้างเคียงไม่ได้ร้องขอไบต์ที่อยู่ข้างเคียงเสมอไป ทำให้ระบบหน่วยความจำต้องให้บริการธุรกรรมที่แยกจากกันมากขึ้น การเปิดใช้งานการโหลดแบบรวมกลุ่ม (coalesced loads) จะทำให้เธรดที่อยู่ติดกันร้องขอไบต์ที่อยู่ติดกัน เพื่อให้ฮาร์ดแวร์สามารถรวมการอ่านของพวกมันเข้าด้วยกันได้

ในการกำหนดค่า bfloat16 การรวมกลุ่ม (coalescing) ช่วยเพิ่มแบนด์วิดท์คีย์จาก 33.8 เป็น 47.9 GB/s, เพิ่มแบนด์วิดท์ค่าจาก 42.0 เป็น 61.8 GB/s และปรับปรุงการ decode แบบ end-to-end ขึ้น 2.1% ที่บริบท 3,840 โทเค็น

แพ็คหัว query เพื่อนำแถว KV กลับมาใช้ซ้ำ

Grouped-query attention ช่วยให้หัวการใส่ใจ (query head) แปดหัวใช้หัว KV ร่วมกันได้ ในการทำ ablation หัว query แต่ละหัวจะทำงานใน simdgroup แยกกัน ดังนั้นทั้งแปดหัวจึงร้องขอแถว KV ที่แคชไว้เดียวกันอย่างอิสระ เคอร์เนลที่ปรับแต่งแล้วจะรวมหัว query สี่หัวเข้าไว้ใน threadgroup เดียว ซึ่งจะโหลดแต่ละแถว KV หนึ่งครั้งและนำกลับมาใช้ใหม่ในการคำนวณ attention สี่ครั้ง Threadgroup ที่สองจะจัดการหัวที่เหลืออีกสี่หัว

เทคนิคนี้ ซึ่งมักเรียกว่าการแพ็ค GQA จะทำการคำนวณแบบเดียวกันและสร้างไบต์เอาต์พุตที่เหมือนกัน ในขณะเดียวกันก็ช่วยลดคำขอ KV อิสระ 8 รายการให้เหลือการโหลดแบบแชร์ 2 รายการ เมื่อเทียบกับการทำ ablation แบบไม่ได้แพ็ค วิธีนี้ช่วยปรับปรุงอัตราการส่งข้อมูล decode แบบ end-to-end ขึ้น 23.8% ที่บริบท 32K โทเค็น

สลับเลย์เอาต์ attention ที่บริบทขนาดยาว

ทุกขั้นตอนการ decode ในเลเยอร์ full-attention จะสแกนแคช KV ที่มีอยู่ เลเยอร์แบบ fixed-block จะแบ่งแคชดังกล่าวออกเป็นส่วนๆ ที่เท่ากัน ซึ่ง GPU สามารถประมวลผลแบบขนานได้ การจัดตารางเวลาเพิ่มเติมของมันไม่คุ้มค่าเมื่อแคชมีขนาดเล็ก แต่เลเยอร์แบบ fixed-block จะช่วยให้งานมีความสมดุลสม่ำเสมอยิ่งขึ้นเมื่อบริบทเติบโตขึ้น

สำหรับโมเดลนี้ รันไทม์จะเก็บเส้นทาง attention ทั่วไปไว้ที่ต่ำกว่า 32K โทเค็น และใช้เส้นทางแบบ fixed-block ที่ 32K โทเค็นขึ้นไป การสลับนี้จะเกิดขึ้นเมื่อหัวแต่ละหัวมี 256 ค่า และหัว query แปดหัวแชร์หัว KV ร่วมกัน รูปแบบอื่นๆ จะยังคงอยู่บนเส้นทางทั่วไป การทำ ablation จะปิดใช้งานการสลับนี้และใช้เส้นทางทั่วไปเสมอ การเปิดใช้งานเส้นทาง fixed-block ช่วยปรับปรุงการ decode แบบ end-to-end ขึ้น 7.7% ที่ 32K, 27.4% ที่ 64K และ 40.2% ที่ 128K

ขีดจำกัดของการปรับแต่งเพิ่มเติม

การเปลี่ยนแปลงบางอย่างช่วยปรับปรุงการทำงานแบบแยกเดี่ยว แต่ไม่ได้ช่วยปรับปรุงการอนุมานแบบ end-to-end

การถอดรหัสแบบคาดเดา (Speculative decoding) ซึ่งใช้โมเดลขนาดเล็กเพื่อเสนอโทเค็นให้โมเดลเต็มตรวจสอบ ทำให้อัตราการ decode แบบ batch-1 ช้าลง 18% การตรวจสอบประมวลผลกลุ่มขนาดสองถึงห้าแถว ซึ่งเป็นรูปร่างที่ไม่มีประสิทธิภาพสำหรับฮาร์ดแวร์นี้ และแถวต่างๆ มักจะเลือกผู้เชี่ยวชาญที่แตกต่างกัน ซึ่งเป็นการเพิ่มปริมาณข้อมูลน้ำหนักผู้เชี่ยวชาญที่มีการอ่าน การลดคำศัพท์เอาต์พุตของตัวร่าง (drafter) ช่วยปรับปรุงอัตราการส่งข้อมูลของตัวร่างขึ้น 4.7–5.1% แต่ไม่ได้ทำให้ลูปการคาดเดาทั้งหมดเร็วขึ้น ผลลัพธ์นี้ขึ้นอยู่กับภาระงาน: การปรับใช้งาน Qwen แบบแบตช์ของเราบน Blackwell ใช้การถอดรหัสแบบคาดเดาภายใต้เงื่อนไขที่แตกต่างกัน

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

การวัดขีดจำกัดของฮาร์ดแวร์ยังแสดงให้เห็นพื้นที่ว่างที่เหลืออยู่น้อยในกระบวนการ prefill และ decode หลัก MoE GEMM และ GEMV สามารถเข้าถึงอัตราการอ่านน้ำหนักที่รักษาไว้ได้เร็วที่สุดถึง 97.9% และ 90.3% สำหรับรูปแบบการเข้าถึงของตน การลบการคำนวณทางคณิตศาสตร์ออกจาก sparse GEMV ทำให้ปริมาณงานเปลี่ยนแปลงไปเพียง 0.2% ซึ่งยืนยันว่าการอ่านน้ำหนักเป็นทรัพยากรจำกัดมากกว่าการคำนวณ ในทำนองเดียวกัน การคูณเมทริกซ์ของ prefill สามารถเข้าถึง 93% ของขีดจำกัดเมทริกซ์ทางทฤษฎีในการทำงานแบบแยกเดี่ยว และ 80–86% ภายในโมเดลที่ผ่านการทดสอบ

ประสิทธิภาพแบบ end-to-end

การเปรียบเทียบแบบ end-to-end จะโหลดไบต์เช็คพอยต์ 4 บิตที่เหมือนกันในเอ็นจิ้นทั้งสองและรันทีละคำขอพารันบน M5 Max 40-core และ 128 GB หนึ่งตัว ภายในแต่ละรอบ เอ็นจิ้นทั้งสองจะทำงานสลับลำดับกันเพื่อลดอคติจากภาระงานเบื้องหลังและอุณหภูมิของชิปที่เปลี่ยนแปลงไป เราเปรียบเทียบกับเส้นทางการสร้างโดยตรงที่เร็วที่สุดของ MLX-LM ไม่ใช่เซิร์ฟเวอร์ ดังนั้นการวัดจึงเน้นไปที่การทำงานของโมเดลมากกว่าโอเวอร์เฮดในการให้บริการ

การกวาดทดสอบครอบคลุมความยาวพรอมต์ 10 ระดับสำหรับ prefill และความยาวบริบท 10 ระดับสำหรับ decode ตั้งแต่ 256 ถึง 128K โทเค็น อัตราการส่งข้อมูล prefill จะเพิ่มขึ้นก่อนในขณะที่เอ็นจิ้นกระจายต้นทุนการตั้งค่าคงที่ไปยังโทเค็นที่มากขึ้น Prefill จะสูงสุดที่พรอมต์ประมาณ 4K โทเค็น จากนั้นจะลดลงเนื่องจากเลเยอร์ full-attention ทั้ง 10 เลเยอร์ทำงานมากขึ้นเมื่อพรอมต์เติบโตขึ้น Decode จะยังคงแบนราบเกือบตลอดที่บริบทสั้นๆ และลดลงเมื่อการอ่านแคช KV ที่เพิ่มขึ้นกลายเป็นสิ่งที่มีนัยสำคัญ เอ็นจิ้นแบบกำหนดเองทำงานได้เร็วกว่าในทุกความยาวที่มีการบันทึก

เนื่องจากการทำงานแบบเฉพาะทางสามารถเปลี่ยนลำดับการทำงานของจุดลอยตัวได้ เราจึงตรวจสอบความสอดคล้องตัวเลขเทียบกับ MLX-LM ด้วย ในการเปรียบเทียบแบบบังคับด้วยครู (teacher-forced comparison) เอ็นจิ้นทั้งสองทำนายโทเค็นถัดไปจากคำนำหน้าอ้างอิงเดียวกันที่ตำแหน่ง 192 ตำแหน่ง ป้องกันไม่ให้ความแตกต่างก่อนหน้าส่งผลกระทบต่ออินพุตภายหลัง ความงุนงง (perplexity) ของ Lily สูงกว่าเพียง 0.04% และมันเลือกโทเค็นอันดับสูงสุดตัวเดียวกันที่ 93.35% ของตำแหน่งที่ทดสอบทั้งหมด

สร้างขึ้นสำหรับแพลตฟอร์มภายในเครื่อง

Apple silicon ไม่ใช่ GPU ของศูนย์ข้อมูลขนาดเล็ก แต่เป็นแพลตฟอร์มการอนุมานภายในเครื่องที่สมบูรณ์ซึ่งมีลักษณะเฉพาะตัวด้านฮาร์ดแวร์และซอฟต์แวร์ หน่วยความจำแบบรวม (unified memory) ช่วยให้โหนดเดียวมีเพดานที่สูงมากในเรื่องปริมาณโมเดลและสถานะที่สามารถรองรับได้ M5 Neural Accelerators รองรับงานเมทริกซ์หนาแน่นในการ prefill และ vector ALU จัดการส่วนที่เหลือที่มีแบนด์วิดท์จำกัดและมีการนำกลับมาใช้ซ้ำต่ำในการ decode

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

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

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