สร้างเอนจินอนุมาน LLM ตั้งแต่ต้นด้วย C++ และ CUDA

• วิธีสร้างเอนจินอนุมาน LLM ด้วย C++ และ CUDA โดยไม่พึ่งไลบรารี
• ช่วยให้เข้าใจสแตกทั้งหมดของการอนุมาน LLM และเห็นผลกระทบของการปรับแต่งแบบต่าง ๆ ต่อความเร็วในการอนุมานได้อย่างชัดเจน
• เป้าหมาย: ทำให้โมเดลสามารถ อนุมานได้รวดเร็วในแบบ single batch บนเซิร์ฟเวอร์ CPU + GPU เดียว และให้ความเร็วประมวลผลโทเคนสูงกว่า llama.cpp

1. ภาพรวมสถาปัตยกรรมและการอนุมานของ LLM

• LLM หลักส่วนใหญ่ใช้สถาปัตยกรรมแบบเดียวกัน คือใช้ทรานส์ฟอร์เมอร์บล็อกหลายบล็อกต่อเนื่องกัน
• การโหลดโมเดลคือการกำหนดคลาสทรานส์ฟอร์เมอร์บล็อกที่ปรับแต่งได้ นำมาประกอบเป็นลำดับ และกำหนดค่าเริ่มต้นด้วยค่าน้ำหนักจาก safetensors
• การอนุมานส่วนใหญ่เกิดขึ้นในแบบ single batch และ "ช่วง decode" กินเวลาการทำงานเกือบทั้งหมด

1.1 ภาพรวมการอนุมาน

• การอนุมานแบ่งเป็นช่วง prefill ที่ส่งโทเคนของพรอมป์ต์เข้าโมเดลเพื่อเติม KV cache และช่วง decode ที่ส่งผ่านโมเดลซ้ำ ๆ เพื่อสร้างโทเคน
  • ช่วง Prefill: ประมวลผลโทเคนของพรอมป์ต์และเริ่มต้น KV cache
  • ช่วง Decode: สร้างทีละหนึ่งโทเคน
• KV cache: เก็บคู่คีย์/แวลูจากก่อนหน้าไว้ เพื่อคำนวณ attention กับบริบทก่อนหน้าได้รวดเร็วขึ้น
• forward pass ของโมเดลจะใช้ตาราง embedding แมป token ID ไปเป็น embedding vector แล้วแปลงสถานะผ่านลำดับของทรานส์ฟอร์เมอร์บล็อก

1.2 คอขวดและเบนช์มาร์ก

• คอขวด: บนฮาร์ดแวร์สมัยใหม่ ข้อจำกัดหลักคือแบนด์วิดท์หน่วยความจำ
  • การสร้างแต่ละโทเคนในการอนุมานโมเดลต้องอ่านทั้งโมเดล ทำให้แบนด์วิดท์หน่วยความจำเป็นข้อจำกัดที่หนักกว่าการคำนวณ
• การควอนไทซ์โมเดลช่วยเพิ่มความเร็วในการอนุมานได้อย่างมีประสิทธิภาพ
• ปริมาณโทเคนสูงสุดตามทฤษฎีขึ้นอยู่กับฮาร์ดแวร์ และสามารถตรวจสอบประสิทธิภาพจริงได้ผ่านเอนจินอินเฟอเรนซ์หลายตัว
• ขีดจำกัดความเร็วตามทฤษฎี:
  • AMD EPYC 7702P: สูงสุด 13.6 tok/s (อิง FP16)
  • RTX 4090: สูงสุด 67.1 tok/s (อิง FP16)
• เบนช์มาร์ก:
  • llama.cpp: CPU 8.7 tok/s, GPU 61 tok/s
  • calm: GPU 66 tok/s

2. การอนุมานบน CPU

• อิมพลีเมนเตชันเริ่มต้นบน CPU เป็นแบบเธรดเดียวและรองรับเฉพาะค่าน้ำหนัก FP32
• สามารถเริ่มทำให้โค้ดขนานด้วยมัลติเธรด และเพิ่มประสิทธิภาพด้วย SIMD

2.1 มัลติเธรด

• ใช้ OpenMP เพื่อทำให้การคูณเมทริกซ์-เวกเตอร์ (matmul) และ multi-head attention ทำงานแบบขนาน เพื่อเพิ่มประสิทธิภาพ
• ผลการปรับแต่ง: ความเร็วดีขึ้นจาก 0.6 tok/s → 4.4 tok/s

2.2 การควอนไทซ์ค่าน้ำหนักและการปรับแต่ง SIMD

• การควอนไทซ์: ควอนไทซ์ค่าน้ำหนัก FP32 เป็น FP16 เพื่อลดการใช้หน่วยความจำลงครึ่งหนึ่งและเพิ่มประสิทธิภาพ
• SIMD: ปรับแต่งด้วย AVX2 เพื่อประมวลผลค่า FP32 ได้พร้อมกัน 8 ค่า
• ผลลัพธ์: ทำได้ 8.4 tok/s

3. การอนุมานบน GPU

• สามารถควอนไทซ์โมเดลเป็น FP16 แล้วโหลดขึ้น RTX 4090 เพื่อเริ่มอิมพลีเมนต์การอนุมานบน GPU
• CUDA ช่วยให้รันฟังก์ชัน C++ (เคอร์เนล) แบบขนานบน GPU ได้

3.1 พอร์ตแบบตรงไปตรงมาด้วย CUDA

• สามารถสร้างแบ็กเอนด์ GPU โดยแปลงงานคำนวณบน CPU เป็น CUDA kernel แบบ 1 ต่อ 1
• CUDA kernel ทำงานแบบอะซิงโครนัส แต่จะรันตามลำดับในสตรีมเดียวกัน
• ปัญหา: ใช้เธรดได้ไม่มีประสิทธิภาพ จึงใช้ทรัพยากร GPU ได้ไม่เต็มที่ → ช้าเพียง 2.9 tok/s

3.2 การคูณเมทริกซ์ (matmul) ที่ดีขึ้น

• การคูณเมทริกซ์ใช้เวลารันจำนวนมากบน CPU และสามารถปรับแต่งด้วย OpenMP ได้
• บน GPU สามารถเพิ่มการใช้เธรดได้ด้วยการให้หนึ่งบล็อกรับผิดชอบหนึ่งแถว
• วิธีปรับแต่ง:
  1. หนึ่งบล็อกประมวลผลหนึ่งแถว และเธรดในบล็อกร่วมมือกันคำนวณ
  2. ใช้ warp reduction
• ผลลัพธ์: ความเร็วเพิ่มเป็น 51.7 tok/s

3.3 การรวมเคอร์เนลและการปรับแต่งเพิ่มเติม

• สามารถเพิ่มประสิทธิภาพได้ด้วยการรวมเคอร์เนล
  • การรวมเคอร์เนล: รวมงานคำนวณต่อเนื่องหลายอย่างไว้ในเคอร์เนลเดียว เพื่อลดการเข้าถึงหน่วยความจำและเวลาในการคำนวณ
• ปรับรูปแบบการเข้าถึงหน่วยความจำและ นำพื้นที่กลับมาใช้ซ้ำ จนทำได้ 56.1 tok/s

3.4 การปรับแต่ง Attention และการจัดการบริบทยาว

• ปัญหา: เมื่อบริบทยาวขึ้น attention kernel กลายเป็นคอขวดด้านประสิทธิภาพ
• ทางแก้:
  1. ปรับแต่งการเข้าถึงหน่วยความจำ: ออกแบบใหม่ให้อ่านบล็อกหน่วยความจำที่ต่อเนื่องกัน
  2. ใช้ shared memory แทน atomicAdd เพื่อแก้ปัญหาค่าทศนิยมที่หายไป
• ผลการปรับแต่ง:
  • บริบทสั้น: 63.8 tok/s (เร็วกว่า 61.0 tok/s ของ llama.cpp)
  • บริบทยาว: ทำได้ 58.8 tok/s

3.5 การควอนไทซ์ KV cache และปัญหาการปรับแต่งของคอมไพเลอร์

• เมื่อควอนไทซ์ KV cache เป็น FP16 กลับเกิดประสิทธิภาพลดลง (เพราะคอมไพเลอร์ปรับแต่งได้ไม่ดีพอ)
• ทางแก้: คลายลูปแบบแมนนวลและใช้ memory prefetching
• ผลลัพธ์: เร็วขึ้นราว 2 เท่าเมื่อเทียบกับ FP32 และยังคงประสิทธิภาพบนบริบทยาวที่ 58.8 tok/s

4. ทิศทางการปรับปรุงในอนาคต

• การปรับแต่ง prompt prefill: ประมวลผลหลายโทเคนพร้อมกันเพื่อลดเวลาสร้างโทเคนแรก
• การรวม Attention kernel: ใช้เทคนิคปรับแต่งแบบ FlashAttention
• การควอนไทซ์ที่สูงขึ้น: ใช้ FP8, INT8, INT4 รวมถึงการควอนไทซ์ activation/cache
• การปรับแต่งเคอร์เนล: นำเทคนิคขั้นสูงมาใช้เพื่อเพิ่มแบนด์วิดท์หน่วยความจำและประสิทธิภาพการคำนวณให้สูงสุด
• การใช้ไลบรารี: ใช้ไลบรารีอย่าง cuDNN และ cuBLAS เพื่อลดเวลาที่ต้องใช้ในการปรับแต่ง

สรุปผลลัพธ์:

• ผ่านการปรับแต่งหลายรูปแบบบน CPU และ GPU จนได้ความเร็ว 63.8 tok/s
• ทำประสิทธิภาพได้ใกล้เคียงหรือดีกว่า llama.cpp และ calm
• สร้างเอนจินอนุมาน LLM ประสิทธิภาพสูงได้โดยใช้เพียง C++ และ CUDA โดยไม่พึ่งไลบรารี

[GN⁺]

ความเห็นจาก Hacker News

• ผู้เขียนดีใจที่บทความในบล็อกของตนได้รับความสนใจ และอยากฟังข้อเสนอแนะ
• ผู้อ่านคนหนึ่งชื่นชมว่าบทความยอดเยี่ยม และสงสัยว่าใช้เวลาเขียนนานแค่ไหน
  • ในฐานะคนที่ทำงานด้าน GPGPU ก็อยากเขียนบทความคล้ายกัน แต่ยังลังเลเพราะไม่แน่ใจว่าจะต้องใช้เวลามากแค่ไหน
• ผู้อ่านอีกคนคิดว่าโค้ดยังไม่ได้ใช้ tensor cores หรือคำสั่ง wgmma
  • อธิบายว่าการเขียนโปรแกรมลักษณะนี้ยาก เพราะต้องจัดการหลายงานพร้อมกัน
  • กล่าวถึงว่าด้วยข้อจำกัดด้านแบนด์วิดท์ อาจไม่จำเป็นต้องมีการคำนวณเพิ่มเติม
  • มองว่าโค้ดในบล็อกน่าจะทำงานได้ดีเมื่อต้องพอร์ตไปยังตัวเร่งความเร็วชนิดอื่น
  • กังวลว่าการใช้ wgmma อาจทำให้ความสามารถในการพอร์ตข้ามเจเนอเรชันของ Nvidia แย่ลง
• ผู้อ่านอีกคนกำลังมองหาเนื้อหา Python ที่คล้ายกัน และอยากแชร์กับทีม
  • ต้องการเนื้อหาที่ครบถ้วนในเชิงแนวคิดและกระชับในสไตล์บทช่วยสอน มากกว่าจะเน้นประสิทธิภาพ
• ผู้ใช้คนหนึ่งอยากเปรียบเทียบเวอร์ชัน Mistral ของตนกับประสิทธิภาพโทเค็นต่อวินาที
  • มีคำแนะนำให้อ้างอิงส่วน quantization ใน README
• มีความเห็นว่า __shfl_down ทุกวันนี้ไม่ค่อยถูกแนะนำแล้ว เพราะปัญหาเรื่องการซิงก์ของ warp