เร่งความเร็ว Sonar ด้วยการเก็งกำไร

การถอดรหัสแบบเก็งกำไร (Speculative decoding) ช่วยเพิ่มความเร็วในการสร้างข้อความของโมเดลภาษาขนาดใหญ่ (LLM) โดยใช้โมเดลร่าง (Draft model) ขนาดเล็กและรวดเร็วเพื่อสร้างตัวเลือกคำตอบที่เสร็จสมบูรณ์ ซึ่งจะถูกตรวจสอบโดยโมเดลเป้าหมาย (Target model) ที่ใหญ่กว่า ภายใต้วิธีการนี้ แทนที่การรันโมเดลเป้าหมายที่มีต้นทุนสูงเพียงครั้งเดียวเพื่อสร้างโทเค็นเดียว จะมีการปล่อยโทเค็นหลายตัวออกมาในขั้นตอนเดียว ในที่นี้เราจะนำเสนอรายละเอียดการใช้งานของการถอดรหัสแบบเก็งกำไรในรูปแบบต่างๆ ที่นำมาใช้ที่ Perplexity เพื่อลดความหน่วงระหว่างโทเค็นบนโมเดล Sonar

การถอดรหัสแบบเก็งกำไร

การถอดรหัสแบบเก็งกำไรใช้ประโยชน์จากโครงสร้างของภาษาธรรมชาติและธรรมชาติแบบ auto-regressive ของทรานสฟอร์เมอร์เพื่อเร่งความเร็วในการสร้างโทเค็น แม้ว่าโมเดลขนาดใหญ่อย่าง Llama-70B จะมีความรู้มากกว่าโมเดลขนาดเล็กอย่าง Llama-1B แต่ในงานที่ง่ายกว่าบางงาน ทั้งสองโมเดลก็มีประสิทธิภาพที่ใกล้เคียงกัน ความทับซ้อนนี้แสดงให้เห็นว่าลำดับข้อความบางลำดับถูกสร้างขึ้นได้ดีกว่าโดยโมเดลที่มีต้นทุนต่ำกว่า ปล่อยปัญหาที่ซับซ้อนให้เป็นหน้าที่ของโมเดลที่ใหญ่กว่า ความท้าทายอยู่ที่การพิจารณาว่าคำตอบใดดีกว่าและกระบวนการสร้างของโมเดลขนาดเล็กมีคุณภาพเท่ากับโมเดลขนาดใหญ่หรือไม่

โชคดีที่ LLM เป็นทรานสฟอร์เมอร์แบบ auto-regressive โดยเมื่อได้รับลำดับของโทเค็น มันจะแสดงการแจกแจงความน่าจะเป็นของโทเค็นถัดไป นอกจากนี้ ค่า logits ที่ได้จากฟีเจอร์ระดับกลางที่เกี่ยวข้องกับโทเค็นในลำดับอินพุตยังบ่งชี้ด้วยว่าโมเดลมีแนวโน้มที่จะออกโทเค็นเหล่านั้นออกมาแน่ชัดเพียงใด คุณสมบัตินี้ช่วยให้เกิดการเก็งกำไร: หากลำดับโทเค็นถูกสร้างขึ้นโดยโมเดลขนาดเล็กโดยเริ่มจากคำนำหน้าอินพุต (Input prefix) มันสามารถถูกนำไปรันผ่านโมเดลขนาดใหญ่เพื่อพิจารณาว่าสอดคล้องกับโมเดลเป้าหมายได้ดีเพียงใด คำนำหน้าแต่ละตัวของตัวเลือกจะถูกให้คะแนนด้วยความน่าจะเป็น และตัวเลือกที่ยาวที่สุดซึ่งอยู่เหนือเกณฑ์การยอมรับจะถูกเลือก นอกจากนี้ โมเดลเป้าหมายยังให้โทเค็นถัดไปมาฟรีๆ อีกด้วย: หากโมเดลร่างสร้างโทเค็น n ตัว โทเค็นสูงสุด n + 1 ตัวจะสามารถถูกปล่อยออกมาได้ในขั้นตอนเดียว

ขยาย

ในขั้นตอนการอนุมาน (Inference) กระบวนการสุ่มตัวอย่างแบบเก็งกำไรสามารถแบ่งออกเป็น 4 ขั้นตอนคร่าวๆ ดังนี้:

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

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

การยอมรับ (Acceptance): โมเดลเป้าหมายจะทำงานบนลำดับของโมเดลร่าง โดยสร้างค่า logits ที่สอดคล้องกับโทเค็นร่างแต่ละตัว ความยาวของลำดับที่ยอมรับได้ยาวที่สุดจะถูกกำหนดขึ้น

การสร้างโมเดลเป้าหมาย (Target generation): เนื่องจากโมเดลเป้าหมายสร้างค่า logits ที่ตำแหน่งที่ไม่ตรงกันหรือที่ส่วนท้ายของลำดับ ค่า logits จึงสอดคล้องกับโทเค็นถัดไป ค่า logits เหล่านี้สามารถนำมาสุ่มตัวอย่างเพื่อให้ได้โทเค็นที่มีความแข็งแกร่งจากโมเดลเป้าหมาย ซึ่งช่วยปิดท้ายลำดับ

มีวิธีการต่างๆ มากมายในการใช้งานการถอดรหัสแบบเก็งกำไร ในบทความนี้ เราจะมุ่งเน้นไปที่วิธีที่เราใช้เพื่อเร่งความเร็วโมเดล Sonar โดยใช้โมเดล 1B ที่พัฒนาขึ้นเองภายในองค์กร รวมถึงกลไกการทำนายที่เรากำลังพัฒนาขึ้นเพื่อเพิ่มความเร็วให้กับโมเดลในระดับของ DeepSeek

เป้าหมาย-ร่าง (Target-Draft)

การถอดรหัสแบบเก็งกำไรสามารถทำได้โดยการจับคู่ LLM ขนาดเล็กที่มีอยู่แล้วเป็นโมเดลร่างเข้ากับโมเดลเป้าหมายเพื่อสร้างลำดับตัวเลือก ในการใช้งานจริง เราได้เร่งความเร็ว Sonar โดยใช้โมเดล Llama-1B ที่ได้รับการปรับแต่งอย่างละเอียด (Fine-tuned) บนชุดข้อมูลชุดเดียวกับเป้าหมาย แม้ว่าวิธีนี้จะไม่จำเป็นต้องฝึกโมเดลร่างตั้งแต่เริ่มต้น แต่โมเดลขนาดเล็กก็ยังคงใช้ความจุแคช KV จำนวนมากและสร้างค่าใช้จ่ายเพิ่มเติมเล็กน้อยในการเติมล่วงหน้า ซึ่งทำให้ TTFT เพิ่มขึ้น

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

ขยาย

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

EAGLE

EAGLE เป็นรูปแบบของ speculative decoding ซึ่งสำรวจลำดับ draft หลายชุดที่สร้างขึ้นผ่านการไล่เรียงแบบต้นไม้ของโทเค็น draft ที่มีแนวโน้มสูง ต้นไม้ที่กำหนดตายตัว (EAGLE) หรือปรับรูปทรงแบบไดนามิก (EAGLE-2) จะถูกสำรวจโดยใช้การรันโทเค็น draft ต่อเนื่องกัน โดยพิจารณา Top-K ของผู้สมัครในแต่ละโหนดแทนการไล่ตามโทเค็นที่มีคะแนนสูงสุดในลำดับเชิงเส้น จากนั้นจะให้คะแนนลำดับเหล่านั้นและเลือกลำดับที่เหมาะสมและยาวที่สุดเพื่อดำเนินต่อ พร้อมทั้งเพิ่มโทเค็นเพิ่มเติมอีกหนึ่งตัวจากเป้าหมาย

ขยาย

เพื่อให้การทำนายมีความแม่นยำยิ่งขึ้น โมเดลร่างของ EAGLE จะไม่เพียงแต่ทำนายตามโทเค็นเท่านั้น แต่ยังใช้ฟีเจอร์ของเป้าหมาย (สถานะซ่อนเร้นของเลเยอร์สุดท้าย หรือ last layer hidden states) ของโมเดลเป้าหมายด้วย ข้อเสียของ EAGLE คือความจำเป็นในการฝึกโมเดลร่างขนาดเล็กแบบกำหนดเองที่มีความแม่นยำเพียงพอที่จะสร้างตัวเลือกที่เหมาะสมภายใต้ขีดจำกัดความหน่วงที่ต่ำ โดยทั่วไป โมเดลร่างจะเป็นเลเยอร์ทรานสฟอร์เมอร์เดี่ยวที่เหมือนกับเลเยอร์ตัวถอดรหัสของโมเดลเดิม ซึ่งเชื่อมโยงอย่างแน่นหนากับเป้าหมายโดยผูกเข้ากับการฝังตัว (Embeddings) และการฉายภาพ lm_head เนื่องจากวิธีนี้ต้องการความจุแคช KV น้อยกว่า EAGLE จึงมีขนาดการใช้หน่วยความจำที่ต่ำกว่า

ในการตรวจสอบลำดับแบบต้นไม้ในโมเดลเป้าหมาย จำเป็นต้องใช้มาสก์ความสนใจ (Attention masks) แบบกำหนดเอง น่าเสียดายที่การใช้มาสก์ความสนใจแบบกำหนดเองสำหรับลำดับทั้งหมดจะทำให้ความสนใจช้าลงอย่างมากสำหรับความยาวอินพุตจริง (มากถึง 50%) ซึ่งทำให้ความเร็วบางส่วนที่ได้จากการเก็งกำไรหมดไป ด้วยเหตุนี้เราจึงยังไม่ได้ปรับใช้การสำรวจต้นไม้แบบเต็มรูปแบบในการผลิต โดยหันมามุ่งเน้นที่กรณีพิเศษของการทำนายโทเค็นเดี่ยวผ่านวิธีแบบ MTP ที่นำเสนอในรายงานทางเทคนิคของ DeepSeek-V3 แทน

MTP

วิธีนี้มีความคล้ายคลึงกับการถอดรหัสแบบร่าง-เป้าหมาย ยกเว้นการใช้สถานะซ่อนเร้นควบคู่ไปกับโทเค็นสำหรับการทำนาย งานที่ต้องทำจะมีมากกว่าเล็กน้อยทั้งในขั้นตอนการเติมล่วงหน้าและขั้นตอนการถอดรหัสเมื่อเทียบกับการเก็งกำไรแบบร่าง-เป้าหมายปกติ โมเดลร่างใช้ทั้งโทเค็นและสถานะซ่อนเร้น: โทเค็น t_{i+1} ถูกสุ่มตัวอย่างจากค่า logits L_i ที่สอดคล้องกับโทเค็น t_i ซึ่งในทางกลับกันได้มาจากสถานะซ่อนเร้น H_i ด้วยเหตุนี้ บัฟเฟอร์โทเค็นอินพุตจึงต้องถูกเลื่อนไปทางซ้ายหนึ่งขั้นตอนเทียบกับเวกเตอร์สถานะซ่อนเร้นที่ส่งออกโดยโมเดลเป้าหมาย รูปภาพด้านล่างแสดงจุดที่สอดคล้องกันซึ่งใช้สำหรับการฝึก รวมถึงการเลื่อนในระหว่างการอนุมาน

ขยาย

กระบวนการถอดรหัสมีความคล้ายคลึงกับการถอดรหัสแบบร่าง-เป้าหมายค่อนข้างมาก ยกเว้นการส่งต่อทั้งสถานะซ่อนเร้นและค่า logits การใช้งานของเราใช้การสุ่มตัวอย่างและการประมวลผล logits ที่เกี่ยวข้องร่วมกันทั้งหมด โดยมีความพิเศษเฉพาะในการเรียกใช้ไปข้างหน้าของโมเดล (Forward invocations) เท่านั้น เมื่อมีการทำนายโทเค็นหลายตัว โมเดลร่างจะใช้สถานะซ่อนเร้นของร่างสำหรับการทำนาย พร้อมทั้งเติมข้อมูลในรายการแคช KV ตามฟีเจอร์ของตนเองด้วย ในระยะยาว วิธีนี้อาจทำให้ความ accuracy ลดลง ต่อมา เมื่อรันโมเดลร่างเพื่อเติมข้อมูลในรายการแคช KV สำหรับการทำนายของเป้าหมาย เราจะรันมันบนลำดับทั้งหมดโดยใช้สถานะซ่อนเร้นของเป้าหมายที่มีความแม่นยำมากกว่าเป็นอินพุต เนื่องจากโมเดลร่างเหล่านี้มีขนาดเล็ก ต้นทุนที่เพิ่มขึ้นในการประมวลผลโทเค็นเพิ่มเติมจึงมีน้อยมากจนแทบไม่น่า A

การฝึกหัว MTP (Training MTP Heads)

เพื่อให้ได้รับประโยชน์จาก MTP เราได้สร้างโครงสร้างพื้นฐานที่จำเป็นในการฝึกหัว MTP ที่แนบมากับโมเดลที่ปรับแต่งแล้วของเราบนชุดข้อมูลของ Perplexity ซึ่งทำงานบนโหนดเดียวที่มีอุปกรณ์ 8xH100 ในเวลาประมาณหนึ่งวัน เราสามารถสร้างหัวสำหรับโมเดลตั้งแต่ Llama-1B ไปจนถึง Llama-70B และ DeepSeek V2-Lite ได้ สำหรับโมเดลขนาดใหญ่ เราอาศัยหัว MTP ที่สร้างขึ้นในระหว่างกระบวนการปรับแต่งอย่างละเอียด

เป้าหมายของการฝึก MTP คือการจับคู่สถานะซ่อนเร้นของโมเดลร่างและค่า logits ที่คาดการณ์จากสถานะซ่อนเร้นของเป้าหมาย ให้เข้ากับค่า logits ของโทเค็นถัดไปและสถานะซ่อนเร้นของเป้าหมาย เนื่องจากกระบวนการอนุมานสำหรับสถานะซ่อนเร้นมีต้นทุนสูง เราจึงคำนวณล่วงหน้าโดยใช้การใช้งานโมเดลเป้าหมายที่ได้รับการปรับแต่งสำหรับการอนุมาน เพื่อนำมาใช้ระหว่างการฝึก อย่างไรก็ตาม เพื่อตรวจสอบความถูกต้องของการใช้งาน MTP ในการอนุมานและให้แน่ใจว่าความแตกต่างทางตัวเลขเนื่องจากการควอนไทเซชัน (Quantization) หรือการปรับแต่งไม่ส่งผลขัดขวางผลลัพธ์ สำหรับการประเมินการสูญเสียการตรวจสอบความถูกต้อง (Validation loss) และความ accuracy เราจึงนำการใช้งานการอนุมานของทั้งโมเดลเป้าหมายและโมเดลร่างกลับมาใช้ใหม่อย่างเต็มรูปแบบ

เมื่อปรับขนาดจากชุดข้อมูล ShareGPT ที่ใช้ในเอกสารต้นฉบับไปเป็นตัวอย่างที่มีขนาดใหญ่ขึ้น เราสังเกตว่าสถาปัตยกรรมหัว MTP ที่สรุปและใช้งานในเอกสาร EAGLE ไม่สามารถฝึกสำหรับโมเดลขนาด 70B ได้ ต่างจาก ShareGPT ที่มีลำดับที่สั้นกว่าจำนวนมาก เราฝึกบนพรอมต์ที่ยาวกว่าอย่างเห็นได้ชัดแต่มีจำนวนน้อยกว่าเล็กน้อย เนื่องจากหัว EAGLE ดั้งเดิมมีโครงสร้างที่แตกต่างจากทรานสฟอร์เมอร์ทั่วไปเล็กน้อย เราจึงได้นำเลเยอร์ RMS Normalization บางเลเยอร์ที่ถูกถอดออกไปกลับมาใส่ใหม่อีกครั้ง เราพบว่าวิธีนี้ไม่เพียงแต่ช่วยให้การฝึกสามารถลู่เข้า (Converge) ได้เท่านั้น แต่ยังช่วยเพิ่มความ accuracy ของหัวขึ้นไปอีกไม่กี่เปอร์เซ็นต์ด้วย

ขยาย

นอกจากเลเยอร์นอร์มจะช่วยอำนวยความสะดวกในการฝึกแล้ว การนำเลเยอร์นอร์มกลับมาใช้ใหม่ยังมีความสมเหตุสมผลทางคณิตศาสตร์อีกด้วย หัว MTP จะนำการฝังตัวและการฉายภาพ logit ของโมเดลเป้าหมายกลับมาใช้ใหม่ เนื่องจากพวกมันอาจมีขนาดใหญ่มาก (ประมาณ 2 GB สำหรับ Llama 70B) ในระหว่างการฝึก ส่วนเหล่านี้จะถูกแช่แข็ง และคาดว่าเลเยอร์ MTP จะเรียนรู้ที่จะฝังคำทำนายลงในพื้นที่เวกเตอร์ (Vector space) เดียวกับที่เลเยอร์การฉายภาพของโมเดลเดิมได้เรียนรู้ในระหว่างการฝึก หากถอดเลเยอร์นอร์มออก MLP เดี่ยวจะถูกคาดหวังให้เรียนรู้ฟังก์ชันเดียวกันกับ MLP ที่ตามด้วยเลเยอร์นอร์ม ซึ่งจะขัดขวางการจับคู่ระหว่างสถานะซ่อนเร้นของโมเดลร่างและโมเดลเป้าหมาย

การอนุมานด้วยการถอดรหัสแบบเก็งกำไร

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

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

รันไทม์การอนุมานที่ Perplexity ถูกสร้างขึ้นรอบๆ FlashInfer ซึ่งเป็นตัวกำหนดเมทาดาทาที่จำเป็นต้องสร้างขึ้นเพื่อกำหนดค่าและกำหนดเวลาเคอร์เนลความสนใจ (Attention kernel) เมื่อกำหนดลำดับอินพุตบางลำดับที่สร้างเป็นแบตช์ สำหรับการเติมล่วงหน้า การถอดรหัส หรือการตรวจสอบ งานฝั่ง CPU จะต้องดำเนินการเพื่อจัดสรรบัฟเฟอร์ระดับกลางและเติมข้อมูลในบัฟเฟอร์คงบางตัวที่ใช้ในความสนใจ งานนี้เป็นงานเพิ่มเติมจากต้นทุนของการกำหนดเวลาแบตช์และการจัดสรรหน้า KV ซึ่งก่อให้เกิดความหน่วงที่ต้องซ่อนไว้เพื่อเพิ่มการใช้งาน GPU ให้สูงสุด

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

กำหนดการแบบร่าง-เป้าหมาย (Draft-Target Schedule)

แม้ว่าจะมีขนาดเล็กกว่าโมเดลเป้าหมาย แต่เมื่อใช้ LLM ทั้งหมดเป็นโมเดลร่าง มันก็ยังคงสร้างความหน่วงที่มากพอสมควรบน GPU ซึ่งช่วยให้มีพื้นที่ว่างเพียงพอสำหรับการซ่อนการดำเนินการฝั่ง CPU ที่มีต้นทุนสูง เนื่องจากโมเดลขนาดเล็กไม่ได้รับประโยชน์จากการขนานเทนเซอร์ (Tensor parallelism) จึงมีความไม่ตรงกันระหว่างจำนวนอันดับที่เป้าหมายและโมเดลร่างถูกแบ่งส่วน (Sharded) ข้ามไป ในการใช้งานของเรา โมเดลร่างจะรันเฉพาะบนอันดับผู้นำ (Leader rank) ของกลุ่ม TP เท่านั้น

ขยาย

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

กำหนดการ MTP สำหรับโทเค็นเดี่ยว (MTP Schedule for a Single Token)

แม้ว่า runtime ยังไม่ได้รองรับการสำรวจต้นไม้ draft แบบสไตล์ Eagle แต่เราได้ใช้งานกรณีพิเศษของรูปแบบนี้ โดยพิจารณาลำดับเชิงเส้นของโทเค็น draft ที่สร้างโดยโมเดลซึ่งมีขนาดเทียบเท่ากับเลเยอร์ decoder ของ transformer เพียงหนึ่งเลเยอร์ รูปแบบนี้สามารถใช้สำหรับการทำนาย draft โดยใช้ weights แบบโอเพนซอร์สของ DeepSeek R1 กรณีย่อยของการทำนายโทเค็นเดียวมีความน่าสนใจ เนื่องจากเลเยอร์ MTP ขนาดใหญ่มีอัตราการยอมรับที่สูงพอที่จะคุ้มกับโอเวอร์เฮดของมัน

การกำหนดเวลาแบบ MTP จะมีความซับซ้อนมากกว่าเล็กน้อย เนื่องจากโมเดลร่างมีความเร็วกว่ามาก ทำให้ซ่อนความหน่วงฝั่ง CPU ได้น้อยลง นอกจากนี้ โมเดลร่างยังถูกแบ่งส่วนไปพร้อมกับโมเดลเป้าหมาย ซึ่งต้องอาศัยการถ่ายโอนหน่วยความจำแบบ shared memory สำหรับข้อมูลแบตช์ การรันจะเริ่มต้นด้วยการถ่ายโอนข้อมูลแบตช์และสุ่มตัวอย่างโทเค็นแรกจากค่า logits ที่ส่งต่อมา เช่นเดียวกับวิธี1ก่อนหน้า ถัดมา โมเดลเป้าหมายจะถูกรันเพื่อตรวจสอบโทเค็น โดยประมวลผลโทเค็น 2 * D ตัว ซึ่ง D คือขนาดแบตช์การถอดรหัส วิธีนี้เหมาะอย่างยิ่งสำหรับการทำไมโครแบตช์ (Micro-batching) ในโมเดล Mixture-of-Experts (MoE) ผ่านการเชื่อมต่อระหว่างกันที่ช้ากว่า เช่น InfiniBand เนื่องจากแบตช์จะแบ่งออกเป็นสองส่วนเท่าๆ กัน สถานะซ่อนเร้นของเป้าหมายจะถูกส่งต่อไปยังการรันโมเดลร่างครั้งถัดไป ในขณะที่ค่า logits จะถูกส่งไปยังตัวสุ่มตัวอย่างเพื่อตรวจสอบความถูกต้อง

ขยาย

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

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

ข้อมูลอ้างอิง

การอนุมานที่รวดเร็วจากทรานสฟอร์เมอร์ผ่านการถอดรหัสแบบเก็งกำไร (Fast Inference from Transformers via Speculative Decoding)

EAGLE: การสุ่มตัวอย่างแบบเก็งกำไรต้องอาศัยการทบทวนความไม่แน่นอนของฟีเจอร์ (EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty)

EAGLE-2: การอนุมานโมเดลภาษาที่เร็วขึ้นด้วยต้นไม้ร่างแบบไดนามิก (EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees)

EAGLE-3: การปรับขนาดการเร่งความเร็วการอนุมานของโมเดลภาษาขนาดใหญ่ผ่านการทดสอบขณะฝึก (EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test)

Medusa: เฟรมเวิร์กเร่งความเร็วการอนุมาน LLM แบบง่ายที่มีหัวถอดรหัสหลายหัว (Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads)

การปรับมาตรฐานเลเยอร์ (Layer Normalization)

FlashInfer: เอนจินความสนใจที่มีประสิทธิภาพและปรับแต่งได้สำหรับการให้บริการอนุมาน LLM (FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving)

รายงานทางเทคนิคของ DeepSeek-V3 (DeepSeek-V3 Technical Report)

การอ้างอิง

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}