Mengompilasi LLM menjadi MegaKernel untuk Mewujudkan Inferensi Berlatensi Rendah

 (zhihaojia.medium.com)
15 poin oleh GN⁺ 2025-06-21 | 1 komentar | Bagikan ke WhatsApp
  • Sebuah compiler telah dikembangkan untuk secara otomatis mengubah inferensi LLM menjadi satu megakernel tunggal
  • Pendekatan MegaKernel (kernel persisten) memungkinkan latensi yang sangat rendah dengan mengintegrasikan komputasi dan komunikasi dalam inferensi LLM sepenuhnya ke dalam satu kernel GPU
  • Ada kendala bahwa karena struktur terdistribusi dari framework ML atau library kernel yang ada, menyatukan seluruh pipeline menjadi satu kernel sangat sulit
  • Mirage Persistent Kernel (MPK) secara otomatis mengubah inferensi LLM multi-GPU menjadi megakernel berperforma tinggi melalui compiler dan sistem runtime
  • MPK mengubah graph komputasi menjadi graph tugas yang lebih terperinci, sehingga memaksimalkan software pipelining dan overlap antara komputasi dan komunikasi
  • Saat MPK diterapkan, latensi pembuatan token berkurang dibandingkan sistem yang ada, dan peningkatan performa menjadi makin besar saat jumlah GPU bertambah

Gambaran umum dan keunggulan pendekatan MegaKernel

  • Dalam inferensi large language model (LLM), salah satu cara efektif untuk mengurangi latensi adalah menggabungkan seluruh proses komputasi dan komunikasi ke dalam satu megakernel (kernel yang konsisten)
  • Dengan pendekatan ini, semua pemrosesan, termasuk operasi per layer pada seluruh model hingga komunikasi antar-GPU, dijalankan tanpa jeda oleh satu kernel GPU
  • Keunggulan utamanya adalah sebagai berikut
    • Menghilangkan overhead peluncuran kernel dengan meniadakan pemanggilan kernel berulang
    • Memungkinkan realisasi software pipelining di seluruh layer
    • Menyembunyikan latensi dengan menjalankan komputasi dan komunikasi secara bersamaan

Keterbatasan yang ada dan kemunculan MPK

  • Framework ML yang ada seperti PyTorch, Triton, TVM pada dasarnya tidak mendukung pembuatan megakernel end-to-end secara otomatis
  • Sistem LLM nyata tersusun dari kombinasi berbagai library kernel seperti NCCL/NVSHMEM (komunikasi), FlashInfer/FlashAttention (attention), CUDA/Triton (operasi kustom), sehingga sulit diintegrasikan menjadi satu kernel tunggal
  • Dengan latar belakang ini, para peneliti dari CMU, UW, Berkeley, NVIDIA, dan Tsinghua mengembangkan Mirage Persistent Kernel (MPK)
    • MPK menggabungkan compiler dan runtime untuk secara otomatis mengubah seluruh pipeline inferensi LLM menjadi megakernel berperforma tinggi

Nilai utama MPK

  • MPK menghilangkan sepenuhnya overhead peluncuran kernel dan memaksimalkan overlap antara komputasi, pemuatan data, dan komunikasi antarlapis untuk mewujudkan inferensi LLM berlatensi sangat rendah
  • Dalam pengujian nyata (prompt 39 token, generasi 512 token, tanpa speculative decoding),
    • Pada lingkungan GPU tunggal NVIDIA A100 40GB, dibandingkan latensi decoding per token 14.5ms milik sistem optimalisasi yang ada seperti vLLM/SGLang, MPK menurunkannya hingga 12.5ms
    • Angka ini mendekati batas bawah teoretis (10ms) berdasarkan bandwidth memori 1.6TB/s dan pemuatan bobot 16GB
  • Dengan mengintegrasikan komputasi dan komunikasi sepenuhnya dalam lingkungan multi-GPU, keunggulan performa MPK menjadi semakin menonjol saat jumlah GPU bertambah

Rincian struktur kerja MPK

Part 1. Compiler – graph operasi LLM → graph tugas

  • Secara umum, operasi LLM direpresentasikan sebagai graph komputasi di mana setiap operasi (misalnya perkalian matriks, attention) atau operasi komunikasi (misalnya all-reduce) menjadi node, dan dependensi data menjadi edge
  • Dalam desain yang ada, pendekatan menjalankan kernel terpisah untuk tiap operator adalah hal yang umum, tetapi ini hanya mencerminkan dependensi pada tingkat kernel, bukan pada unit data yang sebenarnya, sehingga peluang pipelining menjadi terbatas
  • Contoh: jika ada all-reduce setelah perkalian matriks, all-reduce baru mulai dijalankan setelah seluruh perkalian matriks selesai. Padahal secara nyata data dapat dipecah dan dependensi parsialnya dimanfaatkan untuk eksekusi parsial
  • Compiler MPK memecah graph operasi menjadi unit yang lebih halus, lalu secara otomatis mengubahnya menjadi fine-grained task graph yang sesuai dengan unit data sebenarnya
    • Setiap task (persegi panjang) adalah unit komputasi/komunikasi yang dialokasikan ke masing-masing GPU SM
    • Setiap event (lingkaran) adalah titik sinkronisasi antar-task
    • Edge antara task dan event merepresentasikan dependensi data/kontrol secara efisien
  • Berkat graph tugas ini, MPK dapat membuat komputasi dan komunikasi saling tumpang tindih secara parsial maupun paralel dengan lebih baik
  • Mirage kernel superoptimizer juga secara otomatis menghasilkan implementasi CUDA berperforma tinggi yang sesuai untuk tiap task

Part 2. Runtime – menjalankan graph tugas di dalam megakernel

  • Runtime MPK menjalankan graph tugas sepenuhnya hanya di dalam satu kernel GPU (megakernel)
  • Semua SM (Streaming Multiprocessors) pada GPU dibagi secara statis ke dalam peran worker dan scheduler

Worker

  • Setiap worker beroperasi pada level SM dan mengelola queue tugas khusus
  • Dengan cara loop, worker akan
    1. Mengambil task berikutnya dari queue
    2. Menjalankannya (misalnya matmul, attention, transfer data)
    3. Memberi notifikasi ke event saat selesai
    4. Mengulangi proses
  • Dengan demikian, pemanfaatan sumber daya tiap worker dapat dioptimalkan dan operasi antarlapis dapat berjalan secara asinkron

Scheduler

  • Scheduler terdistribusi berjalan pada unit satu warp di tiap SM, dan hingga 4 scheduler dapat berjalan secara bersamaan
  • Setiap scheduler mengelola queue event yang aktif, lalu menugaskan task yang syaratnya sudah terpenuhi ke worker
  • Dengan ini, distribusi task dalam skala besar dapat dilakukan tanpa overhead sinkronisasi terpusat

Metode eksekusi berbasis event

  • Saat task selesai, ia akan menaikkan counter event tertentu. Ketika counter mencapai ambang, event diaktifkan dan dimasukkan ke queue scheduler
  • Scheduler kemudian menjalankan task lanjutan yang memiliki relasi dependensi pada event tersebut
  • Berkat mekanisme ini, fine-grained software pipelining dan overlap komputasi-komunikasi terjadi secara alami
    • Contoh: matmul dari satu layer dan attention dari layer lain dapat berjalan bersamaan
    • Komunikasi all-reduce dapat dimulai segera setelah hasil matmul yang selesai sebagian tersedia
  • Karena seluruh penjadwalan dan perpindahan task terjadi di dalam satu konteks kernel, overhead antar-task sangat rendah, pada level 1–2 mikrodetik (μs)

Arah ke depan

  • Tujuan MPK: mendukung agar developer dapat dengan mudah mengompilasi LLM menjadi megakernel hanya dengan menulis sedikit kode Python (sekitar puluhan baris), lalu memperoleh performa maksimum

  • Arah pengembangan utama

    • Dukungan arsitektur GPU terbaru: misalnya untuk NVIDIA Blackwell, termasuk pendekatan yang dioptimalkan pada level warp
    • Penanganan workload dinamis: riset strategi kompilasi untuk model yang memerlukan control flow dinamis seperti mixture-of-experts (MoE)
    • Penjadwalan task tingkat lanjut: meneliti dan mengejar penerapan kebijakan modern seperti berbasis prioritas dan optimasi throughput

  • MPK menghadirkan titik balik mendasar dalam cara kerja kompilasi dan eksekusi inferensi LLM berbasis GPU, serta mengharapkan perluasan kolaborasi dengan komunitas

Materi tambahan

Bacaan terkait

1 komentar

 
GN⁺ 2025-06-21
Opini Hacker News

  • Menarik melihat pendekatan interpreter on-GPU dari penulis tampak seperti arah masa depan yang sangat menjanjikan. Ada riset lain dengan pendekatan yang hampir sama, jadi saya sarankan melihat tulisan terkait. Model pemrograman dasar CUDA (misalnya peluncuran kernel) sedang diakali untuk paralelisme berbasis tugas yang sangat halus, dan saya sendiri melihat langsung bahwa pendekatan ini meningkatkan pemanfaatan hardware. Jadi saya jadi bertanya-tanya apakah CUDA selama ini justru membatasi kita dalam banyak hal. Saya juga berharap riset penulis ini bisa masuk sebagai backend eksperimental PyTorch. Dan satu catatan kecil: dua paragraf pada bagian pertama hampir identik, jadi mungkin ada typo kecil.

    • Terima kasih atas masukannya, dan benar bahwa proyek MegaKernel dari Stanford juga menghadapi tantangan yang serupa. Namun, MPK mengejar pendekatan di mana pengguna mengekspresikan LLM di level PyTorch, lalu compiler secara otomatis mengubahnya menjadi megakernel yang telah dioptimalkan. Tujuannya adalah membuat pemrograman megakernel menjadi mudah diakses oleh siapa saja. Saya sepenuhnya setuju bahwa CUDA menjadi batasan terutama untuk workload yang sensitif terhadap latensi. Seiring GPU makin besar dan cepat, menulis kernel independen yang tetap bisa memanfaatkan sumber daya hardware secara penuh pada batch kecil menjadi semakin sulit. Kami juga secara aktif mengeksplorasi arah di mana MPK dapat membantu mendukung pembuatan megakernel melalui kolaborasi dengan PyTorch. Terima kasih juga sudah menunjukkan paragraf yang duplikat.

  • Saya sudah cukup lama bekerja dekat dengan vLLM dan SGLang, dan saya yakin proyek ini adalah bentuk ideal dari proyek lanjutan berikutnya. Analisis graph dependensi komputasi, fusion operasi, dan penjadwalan task yang lebih cerdas terasa sangat mengesankan. Selamat untuk timnya.

    • Terima kasih atas umpan balik positifnya. Kami sangat antusias bahwa MPK berpotensi membantu memperluas sistem LLM yang ada, terutama di area serving LLM berlatensi rendah. Kami juga bersemangat mengeksplorasi berbagai kolaborasi dan arah pengembangan ke depan.

  • Saya membaca cepat tulisan dan README github-nya, dan menurut saya ini proyek yang sangat keren. Saya penasaran apakah optimisasi seperti ini bisa diterapkan bukan hanya untuk inferensi tetapi juga untuk tahap training. Saya paham bahwa fusion untuk operasi backward dan komunikasi gradient akan menjadi tantangan. Setahu saya saat ini dynamic workload (misalnya MoE) belum didukung, dan saya ingin menyebut paper terbaru tentang pemrosesan MoE dalam satu kernel, FlashDMoE: Fast Distributed MoE in a Single Kernel.

    • Terima kasih sudah membaca tulisan dan README-nya. Dukungan untuk tahap training memang memungkinkan, tetapi secara umum kernel training cenderung lebih besar sehingga overhead peluncuran kernel tidak terlalu menjadi masalah besar; karena itu inferensi, terutama inferensi berlatensi rendah, mendapat manfaat yang lebih besar. Paper FlashDMoE yang Anda bagikan juga kami baca dengan sangat tertarik, dan dukungan untuk model MoE juga menjadi target berikutnya.

    • Secara pribadi saya agak skeptis terhadap menghabiskan waktu untuk optimisasi training berbasis gradient. Banyak task training di dunia nyata memiliki sifat nilai yang diskret, jadi menurut saya pembelajaran berbasis gradient tidak terlalu mampu menangani hal-hal tersebut dengan baik.

  • Langkah berikutnya tentu mimpi langsung compile ke Verilog lalu beli hardware LLM sendiri di aliexpress.

    • Membagikan tulisan yang memperkenalkan teknologi hardware seperti Chisel. Sebelum kemunculan AI dan GPU, gagasan menerjemahkan software langsung menjadi hardware seperti ini memang pendekatan yang menjanjikan. Perkembangan CPU sedang stagnan, dan hasrat untuk mengoptimalkan lapisan perantara antara software dan hardware terus ada, tetapi kemungkinan besar komputasi paralel bergaya GPU akan tetap menjadi jalur akselerasi utama. CPU untuk tujuan umum pada akhirnya mungkin akan tersisa sebagai otak kecil yang mengelola GPU. Meski begitu, saya memperkirakan pendekatan yang langsung mengubah software menjadi hardware akan sulit menjadi arus utama.

    • Jika arsitektur LLM menjadi stabil dalam 5~10 tahun ke depan, mungkin pemetaan langsung ke hardware bisa menjadi sesuatu yang praktis. Dengan teknologi saat ini, ada kemungkinan model dengan puluhan miliar parameter bisa dimasukkan ke dalam satu wafer hanya dengan menggunakan logic gate ultra-presisi-rendah di sekitar 1.5 bit. Saat presisi naik, jumlah gate meningkat secara eksponensial, jadi untuk saat ini mempertahankan memori bobot dan berbagi unit komputasi masih lebih efisien. Di masa depan, pengembangan LLM ultra-presisi-rendah tampaknya akan menjadi tugas penting.

    • Candaan bahwa biaya training saja sudah tinggi, jadi kalau ditambah biaya masker situasinya makin berat, lalu penilaian yang lebih dingin bahwa startup hardware AI pada dasarnya sudah lama mencoba arah seperti ini.

    • Jika benar-benar ada pendekatan LLM-in-a-box, itu akan sangat menarik. Saya mungkin akan segera punya kesempatan bekerja di lingkungan offline (air-gap), dan solusi seperti itu sepertinya akan sangat berguna.

  • Saya menjalankan langsung kodenya di lingkungan Modal GPU, dan angka peningkatan performa yang diklaim dalam riset ini memang bisa direproduksi dalam praktik. Saya membagikan kode hasil proyek mirage. Pada kombinasi Triton + FlashInfer, latensinya sekitar 19.2ms per token, sedangkan di MPK pada kondisi yang sama turun drastis menjadi 7.7ms.

    • Terima kasih sudah mereproduksi hasilnya secara langsung.

  • Dulu saya pernah ikut kompetisi CUDA kecil. Itu algoritme paralel untuk image atau visi komputer, dan di sana saya mencoba terlihat pintar dengan menyimpan hasil antara ke memori sebagai cache. Setelah hasil kompetisi keluar, saya kaget karena orang lain mengirim kode yang jauh lebih cepat daripada saya. Ternyata alasannya sederhana: mereka tidak menyimpan hasil antara seperti itu, tetapi terus menghitung ulang. Biaya komputasi ternyata jauh lebih murah daripada bolak-balik ke memori. Saya menduga proyek ini mungkin serupa. Saat di-compile menjadi megakernel, batas antar-layer menghilang, jadi berbagi hasil antara berkurang dan jumlah komputasi meningkat, tetapi secara keseluruhan keuntungan besar datang dari berkurangnya bolak-balik ke memori. Mungkin ada sweet spot untuk convolution network, tetapi saya tidak tahu bagaimana megakernel menangani bagian itu.

  • Bahkan sekarang metafora baru untuk LLM terus bermunculan. Saya jadi berpikir apakah kita bisa menganggap LLM seperti transistor. Rasanya sekarang mirip tahap komputer seukuran ruangan yang hanya bisa melakukan perkalian dengan kartu berlubang. Menarik membayangkan apa yang akan terjadi jika kita bisa menjalankan 1 juta query o3-pro secara bersamaan.

  • Proyek ini berasal dari CMU (Carnegie Mellon). Ada juga blog dari Hazy Research di Stanford tentang megakernel, No Bubbles. Menarik melihat persaingan yang aktif di bidang ini. (Tambahan) Ada juga paper yang membahas gambaran lebih besar dari proyek "mirage", tetapi tidak membahas pendekatan megakernel tautan paper

    • Penulis postingan sendiri yang menjawab. Ia setuju bahwa riset dengan Stanford berlangsung secara paralel. Perbedaan utamanya adalah mereka fokus pada compiler untuk pembuatan megakernel otomatis.

    • Menyebut bahwa ThunderKittens dari Hazy Research juga merupakan library yang sangat keren. Akhir-akhir ini banyak upaya dicurahkan untuk formalisasi, pipelining, divide-and-conquer, memaksimalkan efisiensi, serta pengembangan compiler/DSL khusus agar model GPU NVIDIA terbaru bisa dimanfaatkan semaksimal mungkin.

  • Angka performa Qwen 8B akan sangat mengesankan jika tervalidasi. Rasanya lebih praktis dibanding pendekatan megakernel sebelumnya. Pendekatan kernel yang dipertahankan satu per SM seperti ini mengingatkan saya pada Larrabee. Saya jadi penasaran bagaimana dunia akan terlihat sekarang jika alih-alih CUDA, kita menempuh jalur proses-thread-SIMD yang lebih tradisional.

  • Ada ide soal membuat LLM tetap berbasis ASIC murni alih-alih inferensi berbasis software. Apakah ada keuntungan biaya? Mungkinkah tetap menyediakan lapisan yang bisa ditangani atau di-fine-tune lebih lanjut lewat software? Mengingat kita tampaknya sudah hampir mencapai tingkat yang 'cukup bagus', mungkin dalam 2~4 tahun ke depan akan ada keputusan untuk membekukan model ke dalam chip terspesialisasi. Saya penasaran di titik mana tepatnya keuntungan hardware yang sangat terspesialisasi benar-benar mulai terlihat.

    • Pertanyaan lanjutan:
      1. Untuk task tertentu seperti autocomplete, routing kata kunci, dan pengenalan suara, seberapa besar perbedaan latensi dan konsumsi daya antara ASIC dan GPU megakernel, dan apakah itu cukup untuk membenarkan adopsi fungsi tetap di edge device.
      2. ASIC memang sulit dilatih ulang, tetapi mungkinkah ada desain hibrida di mana model dasar ditanam di hardware, sementara hanya modul kecil yang dapat dilatih seperti LoRA dijalankan di koprosesor serbaguna.
      3. Apakah topologi transformer yang tetap cocok untuk penggunaan ulang ruang dalam desain ASIC, atau apakah ukuran model setingkat GPT-3 masih terlalu sulit dijadikan ASIC tanpa pruning atau kuantisasi yang agresif.