LLM’leri MegaKernel’e Derleyerek Düşük Gecikmeli Çıkarım Elde Etmek
(zhihaojia.medium.com)- LLM çıkarımını tek bir megakernel’e otomatik dönüştüren bir derleyici geliştirildi
- MegaKernel(Persistent kernel) yaklaşımı, LLM çıkarımında hesaplama ve iletişimi tamamen tek bir GPU çekirdeğinde birleştirerek çok düşük gecikme sağlıyor
- Mevcut ML framework’leri ve kernel kütüphanelerinin dağıtık yapısı nedeniyle tüm pipeline’ı tek bir kernel’e dönüştürmek oldukça zor
- Mirage Persistent Kernel(MPK), derleyici ve çalışma zamanı sistemi aracılığıyla çoklu GPU LLM çıkarımını otomatik olarak yüksek performanslı bir megakernel’e dönüştürüyor
- MPK, hesaplama grafiğini ince taneli bir görev grafiğine dönüştürerek yazılım pipeline’ını ve hesaplama-iletişim örtüşmesini en üst düzeye çıkarıyor
- MPK uygulandığında mevcut sistemlere kıyasla token üretim gecikmesi azalıyor ve GPU sayısı arttıkça performans kazanımı daha da büyüyor
Genel bakış ve MegaKernel yaklaşımının avantajları
- Büyük dil modeli(LLM) çıkarımında gecikmeyi azaltmanın etkili yollarından biri, tüm hesaplama ve iletişim sürecini tek bir megakernel(tutarlı kernel) içinde birleştirmek
- Bu yaklaşımda modelin tüm katman bazlı işlemleri ve GPU’lar arası iletişim dahil olmak üzere bütün süreçler tek bir GPU kernel’i tarafından kesintisiz yürütülür
- Başlıca avantajlar şöyle
- Tekrarlanan kernel çağrılarını ortadan kaldırarak kernel launch overhead’ini yok eder
- Katmanlar genelinde yazılım pipeline’ı gerçekleştirmeyi mümkün kılar
- Hesaplama ve iletişimi eşzamanlı yürüterek gecikmeyi gizler
Mevcut sınırlamalar ve MPK’nin ortaya çıkışı
- PyTorch, Triton, TVM gibi mevcut ML framework’leri, uçtan uca megakernel’i otomatik üretme yeteneğini özünde desteklemiyor
- Gerçek LLM sistemleri NCCL/NVSHMEM(iletişim), FlashInfer/FlashAttention(attention), CUDA/Triton(özel işlemler) gibi farklı kernel kütüphanelerinin birleşiminden oluştuğu için tek bir kernel altında birleştirilmesi zor
- Bu arka planda CMU, UW, Berkeley, NVIDIA ve Tsinghua araştırmacıları Mirage Persistent Kernel(MPK) geliştirdi
- MPK, derleyici ve runtime’ı birleştirerek tüm LLM çıkarım pipeline’ını otomatik olarak yüksek performanslı bir megakernel’e dönüştürüyor
MPK’nin temel değeri
- MPK, kernel launch overhead’ini tamamen ortadan kaldırıyor ve katmanlar arası hesaplama/veri yükleme/iletişim örtüşmesini en üst düzeye çıkararak ultra düşük gecikmeli bir LLM çıkarım ortamı sağlıyor
- Gerçek testlerde(39 token’lık prompt, 512 token üretimi, speculative decoding kullanılmadan),
- Tek NVIDIA A100 40GB GPU ortamında vLLM/SGLang gibi mevcut optimize sistemlerin token başına decoding gecikmesi olan 14.5ms’ye kıyasla MPK bunu 12.5ms’ye kadar düşürüyor
- Bu değer teorik alt sınıra(10ms) yaklaşıyor(1.6TB/s bellek bant genişliği, 16GB ağırlık yükleme bazında)
- Çoklu GPU ortamında hesaplama ve iletişimin tamamen birleştirilmesi sayesinde, GPU sayısı arttıkça MPK’nin performans üstünlüğü daha da belirgin hale geliyor
MPK’nin çalışma yapısının ayrıntıları
Part 1. Derleyici – LLM hesaplama grafiği → görev grafiği dönüşümü
- Genel olarak LLM hesaplamaları, her bir işlem(ör. matris çarpımı, attention) veya iletişim işlemini(ör. all-reduce) düğüm; veri bağımlılıklarını ise kenar olarak alan bir hesaplama grafiği ile ifade edilir
- Mevcut tasarımda her operatör için ayrı kernel çalıştırma yaklaşımı yaygın olsa da bu, gerçek veri birimi bağımlılıklarını değil yalnızca kernel düzeyi bağımlılıkları yansıttığı için pipeline fırsatlarını sınırlar
- Örnek: Bir matris çarpımının ardından all-reduce geldiğinde, tüm matris çarpımı bitmeden all-reduce başlayamaz. Oysa veriyi parçalayıp kısmi yürütme ve bağımlılık ilişkilerinden yararlanmak mümkündür
- MPK derleyicisi, hesaplama grafiğini ayrıntılandırarak gerçek veri birimine uygun fine-grained task graph’a otomatik olarak dönüştürüyor
- Her görev(dikdörtgen), tek tek GPU SM’lerine atanan hesaplama/iletişim birimidir
- Her olay(daire), görevler arası senkronizasyon noktasıdır
- Görevler ve olaylar arasındaki kenarlar, veri/kontrol bağımlılıklarını verimli biçimde ifade eder
- Bu görev grafiği sayesinde MPK, hesaplama ve iletişimin kısmen ya da paralel biçimde daha fazla örtüşmesini sağlayabiliyor
- Mirage kernel superoptimizer ile her göreve uygun yüksek performanslı CUDA uygulamaları da otomatik üretiliyor
Part 2. Runtime – megakernel içinde görev grafiğinin yürütülmesi
- MPK runtime’ı, görev grafiğini yalnızca GPU üzerindeki tek bir kernel’in(megakernel) içinde tamamen çalıştıran bir yapıya sahip
- GPU’nun tüm SM’leri(Streaming Multiprocessors), statik olarak worker ve scheduler rollerine bölünüyor
Worker
- Her worker, SM düzeyinde çalışır ve kendine ait bir görev kuyruğunu yönetir
- Döngü şeklinde
- Kuyruktan bir sonraki görevi alır
- Yürütür(ör. matmul, attention, veri aktarımı)
- Tamamlanınca olayı bilgilendirir
- Süreci tekrarlar
- Bu sayede her worker’ın kaynak kullanımı optimize edilir ve katmanlar arası asenkron işlem mümkün olur
Scheduler
- Dağıtık scheduler, her SM içinde tek bir warp düzeyinde çalışır ve aynı anda en fazla 4 scheduler yürütülebilir
- Her scheduler, etkinleşen olay kuyruğunu yönetir ve koşulları sağlanan görevleri worker’lara atar
- Böylece merkezi senkronizasyon overhead’i olmadan büyük ölçekli görev dağıtımı yapılabilir
Olay tabanlı yürütme yaklaşımı
- Bir görev tamamlandığında belirli bir olay sayacı artırılır. Sayaç eşik değere ulaştığında olay etkinleşir ve scheduler kuyruğuna eklenir
- Scheduler, bu olaya bağımlı sonraki görevleri çalıştırır
- Böylece ince taneli yazılım pipeline’ı ve hesaplama-iletişim örtüşmesi doğal biçimde gerçekleşir
- Örnek: Bir katmandaki matmul ile başka bir katmandaki attention aynı anda çalışabilir
- Matmul’un kısmen tamamlanan sonuçları çıkar çıkmaz all-reduce iletişimi başlayabilir
- Tüm scheduling ve görev geçişleri tek bir kernel bağlamı içinde gerçekleştiği için görevler arası overhead 1–2 mikrosaniye(μs) düzeyinde kalır
Gelecek yönelimleri
-
MPK’nin hedefi: Geliştiricilerin yalnızca az miktarda Python kodu(yaklaşık birkaç onlarca satır) yazarak LLM’leri kolayca megakernel’e derleyebilmesini ve en yüksek performansı elde edebilmesini sağlamak
-
Başlıca gelişim yönleri
- En yeni GPU mimarisi desteği: Örneğin NVIDIA Blackwell hedefi, warp düzeyi özelleştirilmiş yaklaşımlar vb.
- Dinamik workload işleme: mixture-of-experts(MoE) gibi dinamik kontrol akışı gerektiren modeller için derleme stratejileri araştırmak
- Gelişmiş görev zamanlama: Öncelik tabanlı, throughput optimizasyonlu modern politikaların araştırılması ve uygulanması
-
MPK, GPU tabanlı LLM çıkarım işlerinin derlenme ve çalıştırılma biçiminde köklü bir dönüm noktası öneriyor ve toplulukla iş birliğinin genişlemesini hedefliyor
Ek kaynaklar
- MPK(Mirage Persistent Kernel) kodu ve dokümantasyonu ile en güncel araştırma çıktıları GitHub(https://github.com/mirage-project/mirage) üzerinden incelenebilir
1 yorum
Hacker News yorumu
Yazara: GPU üzerinde çalışan yorumlayıcı yaklaşımının çok umut verici bir gelecek yönü gibi görünmesi ilginç. Neredeyse aynı yaklaşımı gösteren başka bir çalışma da var; bu yüzden ilgili yazıya bakmanızı öneririm. CUDA’nın temel programlama modelinin (ör. kernel launch) ince taneli görev tabanlı paralelleştirme için baypas edildiğini görüyorum ve bu yaklaşımın donanım kullanımını doğrudan artırdığına bizzat tanık oldum. Acaba CUDA bizi pek çok açıdan kısıtlıyor muydu? Yazarın çalışmasının PyTorch’un deneysel bir backend’i olarak yer alma ihtimali heyecan verici. Ayrıca ilk bölümdeki iki paragraf neredeyse aynı; küçük bir yazım hatası notu.
vLLM ve SGLang ile bir süredir yakın çalışıyorum; bu projenin tam olarak bir sonraki ideal devam projesi olduğuna inanıyorum. Hesaplama bağımlılık grafiğini analiz edip işlemleri fuse etme veya görevleri daha akıllıca zamanlama yaklaşımı etkileyici. Ekibi tebrik ederim.
Yazıyı ve github README’sini gözden geçirdim; gerçekten harika bir proje. Bu tür optimizasyonların sadece çıkarımda değil eğitim aşamasında da uygulanıp uygulanamayacağını merak ediyorum. Özellikle backward işlemleri ile gradient iletişiminin fuse edilmesinin zorlu bir konu olduğunun farkındayım. Şu an dynamic workload’ları (ör. MoE) desteklemediğini biliyorum; bu bağlamda yakın zamanda çıkan FlashDMoE: Fast Distributed MoE in a Single Kernel makalesini anmak isterim.
Yazıyı ve README’yi okuduğunuz için teşekkürler. Eğitim aşamasını desteklemek mümkün olsa da, eğitim kernel’leri genelde daha büyük olduğu için kernel launch overhead’i büyük bir sorun olmuyor; bu nedenle en büyük faydayı çıkarım, özellikle de düşük gecikme senaryoları görüyor. Paylaştığınız FlashDMoE makalesini de ilgiyle incelemişler ve MoE model desteğini de bir sonraki hedef olarak gördüklerini vurguluyorlar.
Kişisel olarak gradient tabanlı eğitim optimizasyonuna zaman yatırma konusunda biraz şüpheliyim. Pratikte birçok eğitim görevinin ayrık değer doğası taşıdığını ve gradient tabanlı eğitimin bunları iyi ele alamadığını düşünüyorum.
Bir sonraki adım olarak doğrudan Verilog’a derleyip aliexpress’ten kendi LLM donanımımı satın almak hayalim.
Chisel gibi donanım teknolojilerini tanıtan bir yazı paylaşılıyor. AI ve GPU’lardan önce, yazılımdan doğrudan donanıma geçme fikri umut verici bir yaklaşımdı. CPU gelişimi durağanlaşmış durumda ve yazılım ile donanım arasındaki ara katmanı daha fazla optimize etme arzusu sürse de, GPU tarzı paralel hesaplamanın ana hızlandırma yöntemi olmayı sürdüreceği tahmin ediliyor. Genel amaçlı CPU’lar sonunda GPU’ları yöneten küçük birer beyin rolünde kalabilir. Yine de yazılımdan doğrudan donanıma geçiş yaklaşımının ana akım olması zor görünüyor.
5–10 yıl sonra LLM mimarisi istikrar kazanırsa, bunu doğrudan donanıma eşleştirmenin pratik hale gelebileceği öngörülüyor. Mevcut teknolojiyle on milyarlarca parametrenin bile tek bir wafer’a, 1.5 bit civarında ultra düşük hassasiyetli mantık kapıları kullanılarak sığabileceği belirtiliyor. Hassasiyet arttıkça kapı sayısı geometrik olarak arttığından, bugün için ağırlık belleğini koruyup hesaplama birimlerini paylaşma yaklaşımı daha verimli. Gelecekte ultra düşük hassasiyetli LLM geliştirmek zorunlu bir görev olabilir.
Eğitim maliyetleri zaten yüksekken maske maliyeti de eklenince işlerin daha da zorlaştığına dair bir espri ve aslında AI donanım girişimlerinin uzun süredir böyle yönleri denediğine dair ayık bir değerlendirme.
Eğer LLM-in-a-box yaklaşımı gerçekten varsa oldukça çekici olurdu. Yakında offline (air-gap) bir ortamda çalışma fırsatım olabilir; böyle bir çözüm çok faydalı görünüyor.
Kodu doğrudan Modal GPU ortamında çalıştırdım ve çalışmada iddia edilen performans artışı rakamlarının gerçekten yeniden üretilebildiğini gördüm. mirage proje sonuç kodunu paylaşıyorum. Triton + FlashInfer kombinasyonunda token başına gecikme yaklaşık 19.2ms iken, MPK’de aynı koşullarda 7.7ms’ye düştüğünü deneyimledim.
Eskiden küçük bir CUDA yarışmasına katılmıştım. Görüntü ya da vision alanında paralel bir algoritmaydı; ben de akıllı görünmek için ara sonuçları belleğe cache’lemiştim. Yarışma sonuçlarına baktığımda, başkalarının benden çok daha hızlı kod teslim ettiğini görünce şaşırdım. Sebep şuymuş: ara sonuçları cache’lemek yerine sürekli yeniden hesaplıyorlarmış. Belleğe gidip gelmenin maliyeti, hesaplama maliyetinden çok daha yüksekmiş. Sanırım bu proje de benzer bir noktada duruyor. Megakernel’e derlenince katman sınırları ortadan kalkıyor; ara sonuç paylaşımı azalıyor ve hesaplama miktarı artıyor, ama toplamda belleğe gidiş gelişler azaldığı için büyük kazanç sağlanıyor. Özellikle konvolüsyon ağlarında bir sweet spot olacaktır, ama megakernel’in bunu nasıl ele aldığını bilmiyorum.
LLM için yeni benzetmeler gelmeye devam ediyor. Acaba LLM’leri transistor gibi düşünmek mümkün mü diye aklıma geliyor. Şu anki durum sanki çarpma işlemini sadece delikli kartlarla yapan oda büyüklüğünde bilgisayarlar dönemine benziyor. Aynı anda 1 milyon o3-pro sorgusu çalıştırabilsek neler olacağını hayal etmek eğlenceli.
Bu proje CMU’dan (Carnegie Mellon) çıkmış. Stanford’daki Hazy Research’ün de megakernel üzerine bir blog yazısı var: No Bubbles. Bu alanda rekabetin ne kadar canlı olduğunu görmek etkileyici. (Ek) "mirage" projesinin daha büyük resmini ele alan bir makale de var, ancak megakernel yaklaşımını işlemiyor: makale bağlantısı
Gönderinin yazarı doğrudan yanıt veriyor. Stanford’daki araştırmayla paralel biçimde ilerlediklerini kabul ediyor. Temel farkın, otomatik megakernel üreten bir derleyiciye odaklanmaları olduğunu söylüyor.
Hazy Research’ün ThunderKittens’ının da çok havalı bir kütüphane olduğu belirtiliyor. Son dönemde NVIDIA GPU modellerinden en yüksek verimi almak için formelleştirme, pipelining, böl ve fethet, verimlilik maksimizasyonu ve özel derleyici/DSL geliştirme gibi alanlara çok büyük çaba harcandığı değerlendiriliyor.
Qwen 8B performans rakamları doğrulanırsa oldukça etkileyici olur. Önceki megakernel yaklaşımlarına göre daha pratik hissettiriyor. Her SM başına bir tane olacak şekilde sürekli tutulan bu kernel tarzı, bana eski Larrabee’yi hatırlatıyor. Eğer bugünkü CUDA yerine geleneksel process-thread-SIMD yolu izlenmiş olsaydı dünyanın nasıl olacağını merak ediyorum.
Yazılım tabanlı çıkarım yerine saf ASIC yaklaşımıyla sabit işlevli bir LLM üretme fikri. Maliyet avantajı olur mu? Yazılım tarafında ek katmanlar veya ince ayar yapılabilecek bir yapı sağlanabilir mi? Gerçekten de artık “yeterince iyi” seviyeye çok yaklaşılmış gibi göründüğünden, önümüzdeki 2–4 yıl içinde uzmanlaşmış çiplere sabitleme kararı alınabilir. Aşırı özelleşmiş donanımın avantajının tam olarak hangi noktada belirginleşeceğini merak ediyorum.