Accelerating Sonar through speculation

Speculative decoding mempercepat kecepatan pembuatan Model Bahasa Besar (LLM) dengan menggunakan model draf yang cepat dan kecil untuk menghasilkan kandidat penyelesaian yang diverifikasi oleh model target yang lebih besar. Dalam skema ini, alih-alih menjalankan model target yang mahal untuk menghasilkan satu token, beberapa token dipancarkan dalam satu langkah. Di sini kami menyajikan detail implementasi dari berbagai jenis speculative decoding yang diterapkan di Perplexity untuk mengurangi latensi antar-token pada model Sonar.

Speculative Decoding

Speculative Decoding memanfaatkan struktur bahasa alami dan sifat auto-regresif dari transformer untuk mempercepat pembuatan token. Meskipun model yang lebih besar, seperti Llama-70B, membawa lebih banyak pengetahuan daripada model yang lebih kecil, seperti Llama-1B, pada beberapa tugas yang lebih sederhana kinerjanya serupa. Tumpang tindih ini menunjukkan bahwa urutan tertentu lebih baik dihasilkan oleh model yang lebih murah, sementara masalah yang kompleks diserahkan kepada model yang lebih besar. Tantangannya terletak pada penentuan penyelesaian mana yang lebih baik dan apakah pembuatan dari model yang lebih kecil memiliki kualitas yang sama dengan model yang lebih besar.

Untungnya, LLM adalah transformer auto-regresif: ketika diberi urutan token, mereka mengeluarkan distribusi probabilitas dari token berikutnya. Selain itu, logit yang berasal dari fitur perantara yang terkait dengan token dalam urutan masukan juga menunjukkan seberapa besar kemungkinan model tersebut mengeluarkan token yang tepat tersebut. Properti ini memungkinkan spekulasi: jika urutan token dihasilkan oleh model yang lebih kecil mulai dari prefiks masukan, urutan tersebut dapat dijalankan melalui model yang lebih besar untuk menentukan seberapa baik urutan tersebut selaras dengan model target. Setiap prefiks dari kandidat diberi skor dengan probabilitas dan prefiks terpanjang di atas ambang batas penerimaan dipilih. Sebagai bonus, model target juga menyediakan token berikutnya secara cuma-cuma: jika model draf menghasilkan n token, hingga n + 1 token dapat dipancarkan dalam satu langkah.

Perbesar

Pada waktu inferensi, proses pengambilan sampel spekulatif dapat dibagi secara kasar menjadi 4 tahap:

Prefill: model target dan draf harus dijalankan pada urutan masukan untuk mengisi entri tembolok KV. Meskipun beberapa skema, seperti Medusa, menggunakan lapisan padat yang lebih sederhana untuk prediksi, dalam postingan ini kami berfokus pada draf berbasis transformer yang memerlukan tembolok KV mereka sendiri.

Pembuatan draf: model draf beriterasi untuk menghasilkan sejumlah token tetap. Urutan draf bisa linear atau model dapat mengeksplorasi struktur seperti pohon hingga kedalaman tertentu (EAGLE, Medusa). Di sini, kami berfokus pada urutan linear.

Penerimaan: model target berjalan pada urutan draf, membangun logit yang sesuai dengan setiap token draf. Panjang urutan yang dapat diterima terpanjang ditentukan.

Pembuatan target: karena target menghasilkan logit, pada posisi yang tidak cocok atau bagian akhir urutan, logit tersebut sesuai dengan token berikutnya. Logit ini dapat diambil sampelnya untuk menghasilkan token yang kuat dari model target, menutup urutan tersebut.

Berbagai metode ada untuk mengimplementasikan speculative decoding. Dalam postingan ini, kami akan fokus pada skema yang kami gunakan untuk mempercepat model Sonar menggunakan model 1B internal, serta mekanisme prediksi yang sedang kami bangun untuk mempercepat model pada skala DeepSeek.

Target-Draft

Speculative decoding dapat dicapai dengan menggabungkan 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 (fine-tuned) pada dataset yang sama dengan target. Meskipun pendekatan ini tidak memerlukan pelatihan draf dari awal, model kecil tersebut tetap menggunakan kapasitas tembolok KV yang signifikan dan memperkenalkan sedikit overhead prefill, yang meningkatkan TTFT.

Di bawah skema ini, dekoder hanya berspekulasi pada batch khusus-dekode, menghasilkan token melalui pengambilan sampel standar selama prefill atau pada batch prefill-dekode campuran. Pada tahap prefill, logit target langsung diambil sampelnya untuk melakukan prefill juga pada token yang baru dihasilkan di tembolok KV draf. Draf belum diambil sampelnya, tetapi logit yang dihasilkannya terbawa ke tahap dekode.

Perbesar

Dalam dekode, model draf dimajukan, mengambil sampel token teratas pada setiap tahap. Setelah panjang draf yang diinginkan tercapai, token tersebut dijalankan melalui model target untuk menghasilkan logit yang digunakan sampler untuk mengidentifikasi panjang urutan yang diterima. Penerimaan ditentukan dengan membandingkan distribusi probabilitas penuh dari draf dan target. Karena target selalu menghasilkan satu set logit yang mengikuti urutan draf yang diterima, set tersebut diambil sampelnya untuk menghasilkan output tambahan. Karena model draf belum melihat token yang diterima tersebut, model draf dijalankan ulang untuk mengisi entri tembolok KV yang sesuai sebagai persiapan untuk langkah dekode berikutnya, dengan membawa serta logit.

EAGLE

EAGLE adalah skema speculative decoding yang mengeksplorasi beberapa urutan draf, yang dihasilkan melalui penelusuran seperti pohon dari token draf yang mungkin. Pohon yang tetap (EAGLE) atau berbentuk dinamis (EAGLE-2) dieksplorasi menggunakan eksekusi berturut-turut dari token draf, dengan mempertimbangkan kandidat Top-K di setiap simpul alih-alih mengikuti token dengan skor tertinggi dalam urutan linear. Urutan tersebut kemudian diberi skor dan urutan yang paling sesuai dipilih untuk dilanjutkan, serta menambahkan token tambahan dari target.

Perbesar

Untuk mencapai prediksi yang lebih akurat, model draf EAGLE memprediksi tidak hanya berdasarkan token, tetapi juga menggunakan fitur target (keadaan tersembunyi lapisan terakhir) dari model target. Kerugian dari EAGLE adalah perlunya melatih model draf khusus yang kecil yang cukup akurat untuk menghasilkan kandidat yang sesuai dalam anggaran latensi rendah. Biasanya, model draf adalah satu lapisan transformer yang identik dengan lapisan dekoder dari model asli, yang terikat erat dengan target dengan menghubungkannya ke penyematan (embeddings) dan proyeksi lm_head. Karena ini membutuhkan kapasitas tembolok KV yang lebih sedikit, EAGLE memiliki jejak memori yang lebih rendah.

Untuk memverifikasi urutan seperti pohon dalam model target, masking perhatian khusus harus digunakan. Sayangnya, menggunakan masking perhatian khusus untuk seluruh urutan secara signifikan memperlambat perhatian untuk panjang masukan yang realistis (hingga 50%), meniadakan beberapa percepatan yang dapat dicapai melalui spekulasi. Kami belum menerapkan eksplorasi pohon penuh ke produksi karena alasan ini, dan malah berfokus pada kasus khusus prediksi token tunggal melalui skema mirip MTP yang disajikan dalam Laporan Teknis DeepSeek-V3.

MTP

Skema ini mirip dengan dekode draf-target, dengan pengecualian bahwa keadaan tersembunyi (hidden states) digunakan bersama dengan token untuk prediksi. Sedikit lebih banyak pekerjaan harus dilakukan pada tahap prefill dan dekode dibandingkan dengan spekulasi draf-target reguler. Model draf menggunakan token dan keadaan tersembunyi: token t_{i+1} diambil sampelnya dari logit L_i yang sesuai dengan token t_i, yang pada gilirannya diturunkan dari keadaan tersembunyi H_i. Akibatnya, buffer token masukan harus digeser satu langkah ke kiri relatif terhadap vektor keadaan tersembunyi yang dikeluarkan oleh target. Gambar di bawah ini menandai korespondensi yang digunakan untuk pelatihan, serta pergeseran selama inferensi.

Perbesar

Alur dekode cukup mirip dengan dekode draf-target, dengan pengecualian bahwa keadaan tersembunyi dan logit terbawa. Implementasi kami membagikan semua pemrosesan logit dan pengambilan sampel yang terkait, hanya mengkhususkan pemanggilan penerusan model. Ketika beberapa token diprediksi, model draf menggunakan keadaan tersembunyi draf untuk prediksi, dan juga mengisi entri tembolok KV berdasarkan fiturnya sendiri. Dalam jangka panjang, ini dapat menurunkan akurasi. Selanjutnya, saat menjalankan model draf untuk mengisi entri tembolok KV untuk prediksi target, kami menjalankannya pada seluruh urutan dengan mengambil keadaan tersembunyi target yang lebih akurat sebagai masukan. Karena model draf ini kecil, biaya tambahan untuk memproses token tambahan dapat diabaikan.

Melatih Kepala MTP

Untuk mendapatkan manfaat dari MTP, kami membangun infrastruktur yang diperlukan untuk melatih kepala MTP yang terikat pada model kami yang telah disesuaikan pada dataset Perplexity, yang berjalan pada satu node dengan perangkat 8xH100. Dalam sekitar satu hari, kami dapat membangun kepala untuk model mulai dari Llama-1B hingga Llama-70B dan DeepSeek V2-Lite. Untuk model yang lebih besar, kami mengandalkan kepala MTP yang dibangun selama proses penyesuaian (fine-tuning).

Target pelatihan MTP adalah mencocokkan keadaan tersembunyi draf dan logit yang diekstrapolasi dari keadaan tersembunyi target dengan logit token berikutnya dan keadaan tersembunyi dari target. Karena inferensi untuk keadaan tersembunyi itu mahal, kami pra-menghitungnya menggunakan implementasi model target yang dioptimalkan untuk inferensi, untuk digunakan selama pelatihan. Namun, untuk memvalidasi implementasi MTP inferensi dan memastikan bahwa perbedaan numerik akibat kuantisasi atau pengoptimalan tidak menghambat hasil, untuk kerugian validasi dan estimasi akurasi, kami menggunakan kembali sepenuhnya implementasi inferensi dari model target maupun draf.

Saat menskalakan dari dataset ShareGPT yang digunakan dalam makalah asli ke sampel yang lebih besar, kami memerhatikan 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 yang lebih pendek, kami melatih pada jumlah prompt yang sedikit lebih kecil namun jauh lebih panjang. Karena kepala EAGLE asli sedikit menyimpang dalam strukturnya dari transformer tipikal, kami memperkenalkan kembali beberapa lapisan Normalisasi RMS yang sebelumnya dihapus. Kami menemukan bahwa ini tidak hanya memungkinkan pelatihan untuk konvergen, tetapi juga meningkatkan akurasi kepala sebesar beberapa persen.

Perbesar

Norma lapisan tidak hanya memfasilitasi pelatihan, memperkenalkan kembali norma tersebut juga secara matematis intuitif. Kepala MTP menggunakan kembali penyematan dan proyeksi logit dari model target, karena ukurannya bisa sangat besar (sekitar 2 GB untuk Llama 70B). Selama pelatihan, bagian ini dibekukan dan harapannya adalah lapisan MTP belajar menyematkan prediksi ke dalam ruang vektor yang sama dengan yang dipelajari oleh lapisan proyeksi dari model asli selama pelatihan. Dengan menghilangkan norma tersebut, satu MLP diharapkan mempelajari fungsi yang sama seperti MLP yang diikuti oleh norma, yang menghambat kecocokan antara keadaan tersembunyi dari model draf dan target.

Inferensi dengan Speculative Decoding

Dalam mesin inferensi, untuk menghasilkan token bagi urutan masukan, token tersebut harus terlebih dahulu dikelompokkan ke dalam batch berukuran wajar, kemudian halaman harus dialokasikan dalam tembolok KV untuk token berikutnya. Token masukan dan informasi halaman KV kemudian dikemas ke dalam buffer yang disiarkan ke semua peringkat paralel yang menjalankan model. Terakhir, metadata disalin ke memori GPU dan model dieksekusi untuk menghasilkan logit yang darinya token berikutnya diambil sampelnya.

Tidak seperti implementasi tertentu yang menghubungkan secara longgar server inferensi draf dan target melalui pembungkus yang mengatur permintaan di antara keduanya, pasangan draf-target kami terikat erat dan berjalan melalui pembuatan secara bersamaan. Penjadwalan batch dan alokasi halaman KV dibagikan antara model untuk semua bentuk speculative decoding: ini menyatukan logika yang menjembatani model dengan server inferensi menyeluruh, karena semuanya mengekspos antarmuka yang sama.

Waktu proses inferensi di Perplexity dibentuk di sekitar FlashInfer, yang menentukan metadata yang perlu dibangun untuk mengonfigurasi dan menjadwalkan kernel perhatian. Mengingat beberapa urutan masukan membentuk batch, untuk prefill, dekode, atau verifikasi, pekerjaan sisi CPU harus dilakukan untuk mengalokasikan buffer perantara dan mengisi buffer konstanta tertentu yang digunakan dalam perhatian. Pekerjaan ini merupakan tambahan dari biaya penjadwalan batch dan alokasi halaman KV, yang juga menimbulkan latensi yang harus disembunyikan untuk memaksimalkan pemanfaatan GPU.

Meskipun kami telah memparalelkan sepenuhnya pekerjaan sisi CPU dan sisi GPU untuk inferensi tanpa spekulasi, kami menemukan bahwa keseimbangan CPU-GPU untuk speculative decoding lebih rumit. Tantangan utama muncul dari fakta bahwa jumlah token yang diterima menentukan panjang urutan untuk eksekusi berikutnya, yang memperkenalkan titik sinkronisasi GPU-ke-CPU yang sulit dihindari. Kami bereksperimen dengan berbagai skema penjadwalan untuk menyembunyikan latensi pekerjaan CPU dengan sebaik-baiknya.

Jadwal Draf-Target

Meskipun lebih kecil dari model target, ketika seluruh LLM digunakan sebagai draf, model tersebut tetap memperkenalkan latensi yang cukup besar pada GPU, memberikan sedikit kelonggaran untuk menyembunyikan operasi CPU yang mahal. Karena model kecil tidak mendapatkan manfaat dari paralelisme tensor, ada ketidakcocokan antara jumlah peringkat tempat target dan draf di-shard. Dalam implementasi kami, model draf berjalan hanya pada peringkat pemimpin dari grup TP.

Perbesar

Seperti yang ditunjukkan sebelumnya, langkah dekode membawa logit ke eksekusi berikutnya. Ini memungkinkan kita untuk menumpangtindihkan satu eksekusi model draf dengan pekerjaan penjadwalan batch 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 sampelnya 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 disiapkan secara berurutan dan eksekusinya dimulai untuk mengisi entri tembolok KV-nya dengan token tambahan yang dihasilkan target. Logit yang dihasilkan oleh eksekusi draf ini pada eksekusi saat ini akan digunakan untuk mengambil sampel token draf pertama pada eksekusi berikutnya. Yang terpenting, saat draf berjalan, batch berikutnya dapat dijadwalkan.

Jadwal MTP untuk Token Tunggal

Meskipun waktu proses belum menyediakan eksplorasi pohon draf gaya Eagle, kami menerapkan kasus khusus dari skema ini, dengan mempertimbangkan urutan linear token draf yang dihasilkan oleh model berukuran satu lapisan dekoder transformer. Skema ini dapat digunakan untuk prediksi draf menggunakan bobot open-source dari DeepSeek R1. Sub-kasus memprediksi satu token sangat menarik, karena lapisan MTP yang besar mencapai tingkat penerimaan yang cukup tinggi untuk membenarkan overhead mereka.

Penjadwalan MTP agak lebih kompleks, karena model draf jauh lebih cepat, sehingga menyembunyikan lebih sedikit latensi sisi CPU. Selain itu, draf di-shard bersama dengan model target, yang memerlukan transfer memori bersama untuk informasi batch. Eksekusi dimulai dengan mentransfer info batch dan mengambil sampel token pertama dari logit yang terbawa, mirip dengan skema sebelumnya. Selanjutnya, target dijalankan untuk memvalidasi token, memproses 2 * D token, di mana D adalah ukuran batch dekode. Ini sangat ideal untuk micro-batching dalam model Mixture-of-Experts (MoE) melalui interkoneksi yang lebih lambat seperti InfiniBand, karena batch terbagi rata menjadi dua bagian. Keadaan tersembunyi dari target terbawa ke eksekusi draf berikutnya, sementara logit diteruskan ke sampler untuk verifikasi.

Perbesar

Dengan melakukan sejumlah kecil pekerjaan tambahan pada GPU, kami 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 ulang dengan informasi batch yang sama dengan target, mengisi entri tembolok KV dan membangun logit serta keadaan tersembunyi untuk eksekusi berikutnya, melakukan beberapa pekerjaan redundan pada token yang tidak diterima. Dalam situasi ini, latensi dari pekerjaan yang tidak terpakai hampir tidak dapat diukur karena ukuran model draf yang kecil. Secara paralel dengan eksekusi draf, panjang urutan ditentukan pada CPU dan penjadwalan batch berikutnya dimulai, tanpa harus menunggu pekerjaan GPU selesai.

Overhead pekerjaan tambahan pada lapisan draf tidak terlihat dalam perhatian, namun lapisan MLP lebih bermasalah. Karena instruksi perkalian matriks melakukan padding ke batas 64 di sepanjang dimensi jumlah token, jika penggandaan tidak memerlukan lebih banyak blok secara signifikan, overhead tersebut akan tersembunyi. Untuk urutan draf yang lebih panjang, overhead lebih mahal dan skema yang digunakan untuk model draf-target reguler bekerja lebih baik.

Referensi

Fast Inference from Transformers via Speculative Decoding

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

Layer Normalization

FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving

DeepSeek-V3 Technical Report

Kutipan

@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}
}