Membuat mesin inferensi LLM dari nol dengan C++ dan CUDA

• Cara membangun mesin inferensi LLM menggunakan C++ dan CUDA tanpa library
• Melalui ini, kita bisa memahami seluruh stack inferensi LLM dan merasakan secara langsung bagaimana berbagai optimasi memengaruhi kecepatan inferensi
• Tujuan: mengimplementasikan model agar dapat melakukan inferensi cepat dengan batch tunggal pada satu server CPU + GPU dan mencapai kecepatan pemrosesan token yang lebih cepat daripada llama.cpp

1. Gambaran umum arsitektur dan inferensi LLM

• Sebagian besar LLM utama mengikuti arsitektur yang sama, menggunakan blok transformer yang berurutan.
• Pemuatan model dilakukan dengan mendefinisikan kelas blok transformer yang dapat dikustomisasi, menyusunnya dalam sebuah urutan, lalu menginisialisasinya dengan bobot safetensors.
• Inferensi terutama dilakukan dalam batch tunggal, dan "tahap decode" memakan sebagian besar eksekusi.

1.1 Gambaran umum inferensi

• Inferensi dibagi menjadi tahap prefill, yang mengirim token prompt ke model untuk mengisi cache KV, dan tahap decode, yang berulang kali menjalankan model untuk menghasilkan token
  • Tahap Prefill: memproses token prompt sambil menginisialisasi cache KV
  • Tahap Decode: menghasilkan satu token dalam satu waktu
• Cache KV: menyimpan pasangan key/value sebelumnya agar attention terhadap konteks masa lalu dapat dihitung dengan cepat
• Forward pass model memetakan ID token ke vektor embedding menggunakan tabel embedding, lalu mentransformasikan state melalui urutan blok transformer

1.2 Bottleneck dan benchmark

• Bottleneck: pada hardware modern, bandwidth memori adalah faktor pembatas
  • Saat inferensi model, untuk menghasilkan setiap token perlu membaca seluruh model, sehingga bandwidth memori menjadi kendala yang lebih besar daripada komputasi
• Kuantisasi model efektif untuk meningkatkan kecepatan inferensi
• Throughput token maksimum teoretis berbeda-beda tergantung hardware, dan performa nyata dapat dilihat melalui berbagai mesin inferensi
• Batas kecepatan teoretis:
  • AMD EPYC 7702P: maksimum 13.6 tok/s (berdasarkan FP16)
  • RTX 4090: maksimum 67.1 tok/s (berdasarkan FP16)
• Benchmark:
  • llama.cpp: CPU 8.7 tok/s, GPU 61 tok/s
  • calm: GPU 66 tok/s

2. Inferensi berbasis CPU

• Implementasi awal di CPU berjalan single-thread dan hanya mendukung bobot FP32
• Paralelisasi kode dapat dimulai melalui multithreading, dan performa bisa ditingkatkan menggunakan SIMD

2.1 Multithreading

• Menggunakan OpenMP untuk memparalelkan perkalian matriks-vektor (matmul) dan multi-head attention guna meningkatkan performa
• Hasil optimasi: kecepatan meningkat dari 0.6 tok/s → 4.4 tok/s

2.2 Kuantisasi bobot dan optimasi SIMD

• Kuantisasi: bobot FP32 dikuantisasi ke FP16 sehingga penggunaan memori berkurang setengah dan performa meningkat
• SIMD: dioptimalkan dengan AVX2 agar dapat memproses 8 nilai FP32 sekaligus
• Hasil: mencapai 8.4 tok/s

3. Inferensi berbasis GPU

• Model dapat dikuantisasi ke FP16, dimuat ke RTX 4090, lalu implementasi inferensi GPU dapat dimulai
• Dengan CUDA, fungsi C++ (kernel) dapat dijalankan secara paralel di GPU

3.1 Porting sederhana ke CUDA

• Backend GPU dapat diimplementasikan dengan mengonversi operasi CPU secara 1-1 menjadi kernel CUDA
• Kernel CUDA berjalan secara asinkron, tetapi dalam stream yang sama tetap dieksekusi secara berurutan
• Masalah: karena inefisiensi thread, sumber daya GPU tidak termanfaatkan sepenuhnya → lambat di 2.9 tok/s

3.2 Perkalian matriks (matmul) yang lebih baik

• Perkalian matriks memakan runtime besar di CPU, dan dapat dioptimalkan melalui OpenMP
• Di GPU, utilisasi thread dapat ditingkatkan dengan membuat setiap blok memproses 1 baris
• Metode optimasi:
  1. Satu blok memproses satu baris, dan thread dalam blok berkolaborasi untuk menghitungnya
  2. Menerapkan warp reduction
• Hasil: kecepatan meningkat menjadi 51.7 tok/s

3.3 Fusi kernel dan optimasi tambahan

• Performa dapat ditingkatkan dengan menggabungkan kernel
  • Fusi kernel: menggabungkan operasi yang berurutan ke dalam satu kernel untuk meminimalkan akses memori dan waktu komputasi
• Melalui optimasi pola akses memori dan pemakaian ulang ruang, dicapai 56.1 tok/s

3.4 Optimasi attention dan penanganan konteks panjang

• Masalah: pada konteks panjang, kernel attention menjadi bottleneck performa
• Solusi:
  1. Optimasi akses memori: didesain ulang agar membaca blok memori yang berurutan
  2. Menggunakan memori bersama alih-alih atomicAdd untuk mengatasi masalah nilai floating-point yang terlewat
• Hasil optimasi:
  • Konteks pendek: 63.8 tok/s (lebih cepat daripada 61.0 tok/s milik llama.cpp)
  • Konteks panjang: mencapai 58.8 tok/s

3.5 Kuantisasi cache KV dan masalah optimasi compiler

• Kuantisasi cache KV ke FP16 menyebabkan penurunan performa (karena kurangnya optimasi compiler)
• Solusi: melakukan loop unrolling secara manual dan menerapkan prefetching memori
• Hasil: sekitar 2x lebih cepat dibanding FP32 dan performa konteks panjang tetap 58.8 tok/s

4. Arah peningkatan berikutnya

• Optimasi prefill prompt: memproses beberapa token sekaligus untuk memperpendek waktu hingga token pertama dihasilkan
• Fusi kernel attention: menerapkan teknik optimasi seperti FlashAttention
• Kuantisasi lebih tinggi: menerapkan FP8, INT8, INT4 serta kuantisasi aktivasi/cache
• Optimasi kernel: mengadopsi teknik lanjutan untuk memaksimalkan bandwidth memori dan efisiensi komputasi
• Penggunaan library: memanfaatkan library seperti cuDNN dan cuBLAS untuk mempersingkat waktu optimasi

Ringkasan hasil:

• Melalui berbagai optimasi pada CPU dan GPU, dicapai kecepatan 63.8 tok/s
• Mencatat performa yang mendekati atau lebih baik daripada llama.cpp dan calm
• Mengimplementasikan mesin inferensi LLM berperforma tinggi hanya dengan C++ dan CUDA tanpa library

[GN⁺]

Komentar Hacker News

• Penulis senang karena tulisan blognya mendapat perhatian, dan ingin mendengar masukan
• Seorang pembaca memuji bahwa tulisannya sangat bagus, dan penasaran berapa lama waktu yang dibutuhkan untuk menulisnya
  • Sebagai orang yang bekerja di bidang GPGPU, ia ingin menulis artikel serupa, tetapi ragu karena ketidakpastian waktu yang diperlukan
• Pembaca lain menilai kode tersebut tidak memanfaatkan tensor cores atau instruksi wgmma
  • Ia menjelaskan bahwa pemrograman seperti ini sulit karena harus menangani banyak hal sekaligus
  • Ia menyebutkan bahwa karena batasan bandwidth, mungkin tidak diperlukan komputasi tambahan
  • Ia menilai kode di blog tersebut kemungkinan besar akan bekerja baik saat dipindahkan ke akselerator lain
  • Ia khawatir penggunaan wgmma dapat mengurangi portabilitas antar generasi Nvidia
• Pembaca lain sedang mencari materi Python serupa dan ingin membagikannya kepada timnya
  • Ia menginginkan materi yang ringkas, lengkap secara konseptual, dan bergaya tutorial, alih-alih berfokus pada performa
• Seorang pengguna ingin membandingkan versi Mistral miliknya dan performa token/detik
  • Disarankan untuk melihat bagian kuantisasi di README
• Ada pendapat bahwa __shfl_down belakangan ini tidak direkomendasikan karena masalah sinkronisasi warp