Artikel
Mempercepat Sonar Melalui Spekulasi
Decoding spekulatif mempercepat kecepatan penghasilan Model Bahasa Besar (LLMs) dengan menggunakan model draf yang cepat dan kecil untuk menghasilkan kandidat penyelesaian yang diverifikasi oleh model target yang lebih besar.
Di bawah skema ini, alih-alih menjalankan target mahal yang menghasilkan satu token, beberapa dikeluarkan dalam satu langkah. Di sini kami menyajikan detail implementasi berbagai jenis decoding spekulatif, diterapkan di Perplexity untuk mengurangi latensi antar-token pada model Sonar.
Decoding Spekulatif
Decoding spekulatif memanfaatkan struktur bahasa alami dan sifat auto-regresif transformer untuk mempercepat penghasilan token. Meskipun model yang lebih besar, seperti Llama-70B, memiliki lebih banyak pengetahuan daripada yang lebih kecil, seperti Llama-1B, dalam beberapa tugas yang lebih sederhana mereka menunjukkan kinerja serupa. Tumpang tindih ini menunjukkan bahwa urutan tertentu lebih baik dihasilkan oleh model yang lebih murah, meninggalkan masalah kompleks untuk yang lebih besar. Tantangannya terletak pada menentukan penyelesaian mana yang lebih baik dan apakah penghasilan model yang lebih kecil sama kualitasnya dengan yang lebih besar.
Untungnya, LLM adalah transformer auto-regresif: ketika diberikan urutan token, mereka menghasilkan distribusi probabilitas dari token berikutnya. Selain itu, logit yang berasal dari fitur intermediat yang terkait dengan token dalam urutan masukan juga menunjukkan seberapa mungkin model untuk mengeluarkan token tersebut. Properti ini memungkinkan spekulasi: jika serangkaian token dihasilkan oleh yang lebih kecil mulai dari prefiks masukan, itu dapat dijalankan melalui yang lebih besar untuk menentukan seberapa baik itu cocok dengan model target. Setiap prefiks dari kandidat diberi skor dengan probabilitas dan yang terpanjang di atas ambang penerimaan dipilih. Sebagai bonus, model target juga menyediakan token berikutnya secara gratis: jika model draf menghasilkan n token, hingga n + 1 dapat dikeluarkan dalam satu langkah.
Di waktu inferensi, proses pengambilan spekulatif dapat dibagi menjadi kurang lebih 4 tahap:
Prefill: baik model target dan draf harus dijalankan pada urutan masukan untuk mengisi entri cache KV. Meskipun beberapa skema, seperti Medusa, menggunakan lapisan padat yang lebih sederhana untuk prediksi, dalam pos ini kami fokus pada draf berbasis transformer yang membutuhkan cache KV mereka sendiri.
Penghasilan draf: model draf mengiterasi untuk menghasilkan sejumlah token tetap. Urutan draf bisa linier atau model dapat menjelajahi struktur seperti pohon hingga kedalaman tertentu (EAGLE, Medusa). Di sini, kami fokus pada urutan linier.
Penerimaan: model target dijalankan pada urutan draf, membangun logit yang sesuai dengan setiap token draf. Panjang urutan yang paling dapat diterima ditentukan.
Penghasilan target: sejak target menghasilkan logit, di posisi yang tidak cocok atau ujung ekor urutan, logit sesuai dengan token berikutnya. Logit ini dapat diambil untuk menyediakan token yang kuat dari target, menyempurnakan urutan.
Berbagai metode ada untuk mengimplementasikan decoding spekulatif. Dalam pos ini, kami akan fokus pada skema yang kami gunakan untuk mempercepat model Sonar menggunakan model internal 1B, serta mekanisme prediksi yang kami kembangkan untuk mempercepat model dalam skala DeepSeek.
Target-Draft
Decoding spekulatif dapat dicapai dengan menghubungkan LLM kecil yang ada sebagai model draf ke model target untuk menghasilkan urutan kandidat. Dalam produksi, kami telah mempercepat Sonar menggunakan model Llama-1B yang disesuaikan pada dataset yang sama dengan target. Meskipun pendekatan ini tidak memerlukan pelatihan draf dari awal, model kecil masih menggunakan kapasitas cache KV yang signifikan dan memperkenalkan sedikit overhead prapemenuhan, meningkatkan TTFT.
Di bawah skema ini, decoder hanya berspekulasi pada batch hanya-decode, menghasilkan token melalui pengambilan standar selama prapemenuhan atau pada batch campuran prapemenuhan-decode. Pada tahap prapemenuhan, logit target segera diambil untuk juga mengisi token yang baru dihasilkan dalam cache KV dari draf. Draf belum diambil, tetapi logit yang dihasilkannya dibawa ke tahap decode.
Dalam decode, model draf dimajukan, mengambil token teratas di setiap tahap. Setelah panjang draf yang diinginkan tercapai, token dijalankan melalui model target untuk menghasilkan logit berdasarkan mana sampler mengidentifikasi panjang urutan yang diterima. Penerimaan ditentukan dengan membandingkan distribusi probabilitas penuh dari draf dan target. Karena target selalu mengeluarkan satu set logit setelah urutan draf yang diterima, yang diambil untuk menghasilkan output tambahan. Karena model draf belum melihat token yang diterima itu, model itu dijalankan kembali untuk mengisi entri cache KV yang sesuai sebagai persiapan untuk langkah decode berikutnya, membawa logit lagi.
EAGLE
EAGLE adalah skema decoding spekulatif yang menjelajahi beberapa urutan draf, dihasilkan melalui penelusuran seperti pohon dari token draf yang mungkin. Pohon tetap (EAGLE) atau berbentuk dinamis (EAGLE-2) dieksplorasi menggunakan eksekusi berturut-turut dari token draf, mempertimbangkan kandidat Top-K di setiap node alih-alih mengikuti token dengan skor tertinggi dalam urutan linier. Urutan-urutan itu kemudian diberikan skor dan yang paling sesuai dipilih untuk melanjutkan, juga menambahkan token tambahan dari target.
Dalam rangka mencapai prediksi yang lebih akurat, model draf EAGLE memprediksi tidak hanya berdasarkan token, tetapi juga menggunakan fitur target (layer hidden states terakhir) dari model target. Kerugian dari EAGLE adalah kebutuhan untuk melatih model draf kecil khusus yang cukup akurat untuk menghasilkan kandidat yang sesuai dalam anggaran latensi rendah. Biasanya, model draf adalah satu lapisan transformer yang identik dengan lapisan decoder model asli, yang sangat terhubung dengan target dengan mengikat ke proyeksi embeddings dan lm_head. Karena ini memerlukan kapasitas cache KV yang lebih sedikit, EAGLE memiliki jejak memori yang lebih rendah.
Untuk memverifikasi urutan seperti pohon di model target, masker perhatian khusus harus digunakan. Sayangnya, menggunakan masker perhatian khusus untuk seluruh urutan secara signifikan memperlambat perhatian untuk panjang masukan realistis (hingga 50%), meniadakan beberapa percepatan yang dapat dicapai melalui spekulasi. Kami belum menerapkan penjelajahan pohon penuh ke produksi karena alasan ini, berfokus pada kasus khusus prediksi token tunggal melalui skema seperti MTP yang dipresentasikan dalam Laporan Teknis DeepSeek-V3.
MTP
Skema ini mirip dengan decoding draf-target, dengan pengecualian dari hidden states yang digunakan bersama token untuk prediksi. Sedikit lebih banyak pekerjaan harus dilakukan baik pada tahap prapemenuhan maupun decode dibandingkan dengan spekulasi draf-target biasa. Model draf menggunakan baik token maupun hidden states: token t_{i+1} diambil dari logit L_i yang sesuai dengan token t_i, yang pada gilirannya berasal dari hidden states H_i. Dengan demikian, buffer token masukan harus digeser satu langkah ke kiri relatif terhadap vektor hidden state yang dihasilkan oleh target. Gambar di bawah ini menandai korespondensi yang digunakan untuk pelatihan, serta pergeseran selama inferensi.
Alur decoding sangat mirip dengan decoding draf-target, dengan pengecualian dari hidden states dan logit yang dibawa. Implementasi kami berbagi semua sampling terkait dan logit pemrosesan, hanya mengkhususkan pemanggilan maju model. Ketika beberapa token diprediksi, model draf menggunakan hidden states draf untuk prediksi, juga mengisi entri cache KV berdasarkan fitur-fiturnya sendiri. Dalam jangka panjang, ini dapat menurunkan akurasi. Selanjutnya, ketika menjalankan model draf untuk mengisi entri cache KV untuk prediksi target, kami menjalankannya pada seluruh urutan mengambil hidden states target yang lebih akurat sebagai masukan. Karena model-model draf ini kecil, biaya tambahan untuk memproses token tambahan tidak signifikan.
Pelatihan MTP Heads
Untuk mendapatkan manfaat dari MTP, kami membangun infrastruktur yang diperlukan untuk melatih kepala MTP yang terpasang pada model yang disesuaikan kami pada dataset Perplexity, berjalan pada satu node dengan perangkat 8xH100. Dalam waktu sekitar satu hari, kami dapat membangun kepala untuk model dari Llama-1B hingga Llama-70B dan DeepSeek V2-Lite. Untuk model yang lebih besar, kami mengandalkan kepala MTP yang dibangun selama proses penyesuaian.
Tujuan dari pelatihan MTP adalah untuk mencocokkan hidden states draft dan logit yang diekstrapolasi dari hidden states target ke logit token berikutnya dan hidden states dari target. Karena inferensi untuk hidden states mahal, kami menghitungnya sebelumnya menggunakan implementasi yang dioptimalkan untuk inferensi dari model target, untuk digunakan selama pelatihan. Namun, untuk memvalidasi implementasi MTP inferensi dan memastikan bahwa perbedaan numerik karena kuantisasi atau optimalisasi tidak menghalangi hasil, untuk estimasi kerugian validasi dan akurasi kami sepenuhnya menggunakan kembali implementasi inferensi dari model target dan draf.
Ketika menskalakan dari dataset ShareGPT yang digunakan dalam makalah asli ke sampel yang lebih besar, kami memperhatikan bahwa arsitektur kepala MTP yang diuraikan dan diimplementasikan dalam makalah EAGLE gagal dilatih untuk model berukuran 70B. Tidak seperti ShareGPT yang berisi sejumlah besar urutan pendek, kami melatih pada jumlah sedikit lebih kecil dari prompt yang jauh lebih panjang. Karena kepala EAGLE asli sedikit menyimpang dalam struktur dari transformer tipikal, kami memperkenalkan kembali beberapa lapisan RMS Normalization yang telah dihapus. Kami menemukan bahwa ini tidak hanya memungkinkan pelatihan untuk konvergen, tetapi juga meningkatkan akurasi kepala oleh beberapa poin persentase.
Tidak hanya norma lapisan memfasilitasi pelatihan, mempekenalkan kembali norma juga secara matematis intuitif. Kepala MTP menggunakan kembali embeddings dan proyeksi logit dari model target, karena ini dapat substansial dalam ukuran (sekitar 2 GB untuk Llama 70B). Selama pelatihan, ini dibekukan dan ekspektasi adalah bahwa lapisan MTP belajar memasukkan prediksi ke vektor spasial yang sama dengan yang dipelajari selama pelatihan oleh lapisan proyeksi model asli. Dengan menjatuhkan norma, diharapkan satu MLP belajar fungsi yang sama seperti MLP diikuti oleh norma, yang menghalangi pertandingan antara hidden states model draf dan target.
Inferensi dengan Decoding Spekulatif
Dalam mesin inferensi, untuk menghasilkan token untuk urutan masukan, mereka perlu dikelompokkan terlebih dahulu menjadi batch berukuran wajar, kemudian halaman harus dialokasikan dalam cache KV untuk token berikutnya. Token masukan dan informasi halaman KV kemudian dimasukkan ke dalam buffer yang disiarkan ke semua peringkat paralel yang menjalankan model. Akhirnya, metadata disalin ke memori GPU dan model dijalankan untuk menghasilkan logit dari mana token berikutnya diambil.
Tidak seperti implementasi tertentu yang menghubungkan longgar server inferensi draf dan target melalui pembungkus yang mengatur permintaan di antara mereka, pasangan draf-target kami terhubung erat dan melangkah melalui penghasilan secara serempak. Penyusunan batch dan alokasi halaman KV dibagi di antara model untuk semua bentuk decoding spekulatif: ini menyatukan logika yang menjembatani model dengan server inferensi keseluruhan, karena semuanya mengekspos antarmuka yang sama.
Runtime inferensi di Perplexity dibentuk di sekitar FlashInfer, yang menentukan metadata yang perlu dibangun untuk mengkonfigurasi dan menjadwalkan kernel perhatian. Mengingat beberapa urutan masukan yang membentuk batch, untuk prapemenuhan, decode atau verifikasi, pekerjaan di sisi CPU harus dilakukan untuk mengalokasikan buffer intermedia dan mengisi buffer konstan tertentu yang digunakan dalam perhatian. Pekerjaan ini selain biaya penyusunan batch dan alokasi halaman KV, yang juga menimbulkan latensi yang harus disembunyikan untuk memaksimalkan pemanfaatan GPU.
Sementara kami sepenuhnya memparalelkan pekerjaan di sisi CPU dan GPU untuk inferensi tanpa spekulasi, kami menemukan bahwa keseimbangan CPU-GPU untuk decoding spekulatif lebih rumit. Tantangan utama muncul dari fakta bahwa jumlah token yang diterima menentukan panjang urutan untuk menjalankan berikutnya, memperkenalkan titik sinkronisasi GPU-ke-CPU yang sulit dihindari. Kami bereksperimen dengan berbagai skema penyusunan untuk menyembunyikan latensi pekerjaan CPU dengan sebaik mungkin.
Jadwal Draft-Target
Meskipun lebih kecil dari model target, ketika seluruh LLM digunakan sebagai draf, itu tetap memperkenalkan latensi yang cukup besar pada GPU, memberikan beberapa ruang untuk menyembunyikan operasi CPU yang mahal. Karena model yang lebih kecil tidak mendapatkan manfaat dari paralelisme tensor, ada ketidakcocokan antara jumlah peringkat target dan draf yang terpecah. Dalam implementasi kami, model draf hanya berjalan pada peringkat pemimpin dari kelompok TP.
Seperti yang sudah diindikasikan sebelumnya, langkah decode membawa logit ke dalam putaran berikutnya. Ini memungkinkan kita untuk tumpang tindih satu eksekusi model draf dengan pekerjaan penyusunan batch di sisi CPU. Setelah batch disatukan, panggilan berulang ke sampler dan draf menghasilkan token draf. Secara paralel, batch untuk verifikasi disatukan untuk model target dan disinkronkan dengan pekerja paralel. Logit target diverifikasi dan diambil untuk menentukan panjang urutan yang diterima. Pada titik ini, sinkronisasi GPU-ke-CPU diperlukan untuk menentukan panjang urutan berikutnya. Karena model draf hanya dijalankan pada node pemimpin, batch nya diatur secara berurutan dan eksekusinya dimulai untuk mengisi entri cache KV dengan token tambahan yang dihasilkan target. Logit yang dihasilkan oleh putaran draf ini dalam putaran saat ini akan digunakan untuk mengambil token draf pertama dalam putaran berikutnya. Yang paling penting, saat draf dijalankan, batch berikutnya dapat disusun.
Jadwal MTP untuk Satu Token
Sementara runtime belum menyediakan eksplorasi pohon draf gaya-Eagle, kami mengimplementasikan kasus khusus dari skema ini, mempertimbangkan urutan linier dari token draf yang dihasilkan oleh model dengan ukuran satu lapisan decoder transformer. Skema ini dapat digunakan untuk prediksi draf menggunakan bobot sumber terbuka dari DeepSeek R1. Sub-kasus prediksi satu token menarik, mengingat lapisan MTP besar mencapai tingkat penerimaan yang cukup tinggi untuk membenarkan overhead mereka.
Penyusunan MTP agak lebih kompleks, karena model draf jauh lebih cepat, menyembunyikan lebih sedikit latensi di sisi CPU. Selain itu, draf dipartisi bersama model target, memerlukan transfer memori bersama untuk informasi batch. Satu putaran dimulai dengan mentransfer info batch dan mengambil token pertama dari logit yang dibawa, sama seperti skema sebelumnya. Selanjutnya, target dijalankan untuk memvalidasi token, memproses 2 * D token, di mana D adalah ukuran batch decode. Ini ideal untuk pembuatan batch mikro pada model Mixture-of-Experts (MoE) melalui interkoneksi yang lebih lambat seperti InfiniBand, karena batch terbagi rata menjadi dua bagian. Hidden states target dibawa ke putaran draf berikutnya, sementara logit diteruskan ke sampler untuk verifikasi.
Dengan melakukan sejumlah pekerjaan tambahan yang terbatas pada GPU, kita menghindari sinkronisasi CPU-ke-GPU setelah penerimaan urutan draf. Setelah token masukan dari target digeser, kernel memasukkan token target berikutnya ke lokasi yang sesuai. Draf kemudian dijalankan kembali dengan informasi batch yang sama seperti target, mengisi entri cache KV dan membangun logit dan hidden states untuk putaran berikutnya, melakukan beberapa pekerjaan berlebih pada token yang tidak diterima. Dalam situasi ini, latensi pekerjaan yang tidak terpakai hampir tidak dapat diukur karena kecilnya ukuran model draf. Secara paralel dengan putaran draf, panjang urutan ditentukan di CPU dan penyusunan batch berikutnya dimulai, tanpa harus menunggu pekerjaan GPU selesai.
Overhead dari pekerjaan tambahan di lapisan draf tidak terlihat dalam perhatian, namun lapisan MLP lebih problematik. Karena instruksi perkalian matriks mengisi batas 64 sepanjang dimensi jumlah token, jika menggandakan tidak memerlukan blok yang secara signifikan lebih banyak, overhead tersembunyi. Untuk urutan draf yang lebih panjang, overhead lebih mahal dan skema yang digunakan untuk model draf-target umum bekerja lebih baik.
Referensi
EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty
EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees
EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test
Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads
FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving