Perplexity'de CuTeDSL

Perplexity'de, gömme işlemlerinden trilyon parametreli LLM'lere kadar çeşitli modellere hizmet vermek için dahili Runtime-Optimized Serving Engine (ROSE) kullanıyoruz. Sonar, Search ve Embeddings dahil tüm API'lerin arkasında yer alan çalışma zamanımız; araştırmacıların en uygun modelleri ve mimarileri keşfetmesini sağlarken en iyi adayları da hızla üretime alabilmesi için geliştirilmiştir. Bunu başarmak için, en son teknoloji mimari çıkarımlarını hızlı bir şekilde uygulayabilmemiz, dağıtabilmemiz ve optimize edebilmemiz gerekir. Bu blog yazısında, çıkarım motorumuza entegre edilen CuTeDSL çekirdeklerinin, modellerimizi NVIDIA GPU'larda en yüksek performansa ulaştırmak için gereken GPU çekirdeklerini oluşturmamıza nasıl olanak tanıdığını inceliyoruz.

Perplexity'deki çeşitli ekipler tarafından oluşturulan özel modeller, AI Çıkarımı ekibi tarafından oluşturulan ROSE çıkarım motoru kullanılarak NVIDIA Hopper ve Blackwell GPU'larında dahili olarak barındırılır. ROSE başlangıçta, NVIDIA Triton Inference Server ile uyumlu bir arayüze sahip olarak hem dil kod çözme hem de sınıflandırma için özelleştirilmiş Llama modellerine daha iyi hizmet vermek üzere oluşturulmuştur. Zamanla, Perplexity'nin ürünlerini desteklemek üzere çok çeşitli ekipler için önemli LLM'lerin yanı sıra sıralama, sınıflandırma, puanlama ve diğer transformatör tabanlı modelleri sunan tam teşekküllü bir çıkarım motoruna dönüştü. ROSE'un tasarım hedefi basittir: zaman içinde en yüksek performansa ulaşmak için çaba yatırımı yapılır, ancak asıl öncelik modellerimiz için sürekli gelişen mimarilere ve kullanım durumlarına uyum sağlamaktır.

Veri türleri, baş boyutları, derleme zamanı sabitleri ve başlatma ayarları genelinde çekirdek yapılandırma alanı.

Üst düzey bir bakış açısıyla ROSE, bir modeli istemcilerin etkileşime girebileceği bir arayüze uyarlayan bir motor'den oluşur. Motorlar yüksek düzeyde yeniden kullanılabilir olup, istek zamanlama ve toplu işlem yönetimini gerçekleştirir ve temel modellerden kodlanmış belirteçleri, gömmeleri veya puanları alır. Cihazların başlatılması, süreçler arası ve düğümler arası iletişim kanalları ve ağırlık yükleme işlemlerinin tümü bunlar tarafından yönetilir. LLM'ler için motor; toplu işleme, parçalama (chunking) ve örneklemeden sorumludur. Ayrıca hem tam hem de doğrusal dikkat (attention) için KV depolama alanının tahsisinden ve önek (prefix) eşleştirmeden sorumludur. Gömme motorları, yalnızca çevrimiçi toplu işleme gerçekleştirdikleri için daha basittir. Yürütme, motor, hızlandırıcı üzerinde çalışan model ile paralel olarak çalışırken bir sonraki iş turunu hazırlayacak şekilde optimize edilmiştir.

ROSE, CuTeDSL, Triton, CUDA, CUTLASS ve cuBLAS ile oluşturulmuş GPU çekirdeklerini sarmalayan, modelleme için zengin bir özel katman koleksiyonu sunar. Matris çarpımı ve dikkat (attention) uygulamaları için öncelikle CUTLASS ve cuBLAS'tan NVIDIA çekirdeklerine güveniyoruz. Gömme (embedding), norm, MoE ve aktivasyon çekirdekleri dahil boru hattının geri kalanı, ROSE'un çekirdek kitaplığının bir parçasını oluşturur. Çıkarım çalıştırdığımız donanımı en iyi şekilde desteklemek için ROSE, MoE yönlendirmesi, dağıtımı ve birleştirmesi için kendi cihazlar arası iletişim ilkelleriyle donatılmıştır. Diğer çekirdekler nispeten basit olsa da, üzerinde çalıştıkları tüm farklı gizli boyutlarda ve cihazlarda en yüksek performansı elde etmek için kapsamlı bir şekilde özelleştirilmeleri gerekir. Ek olarak, çeşitli modellerin tüm ayarlarını ve ince ayarlarını barındırmak için derleme zamanı uzmanlaşmasına güveniyoruz (RMS norm için ağırlık sapması ve ağırlık sapması olmaması, aktivasyon fonksiyonları vb.). Sonuç olarak, bunları uygulamak için DSL'leri kullanmaya karar verdik: başlangıçta ROSE'u Triton kullanarak oluşturduk. Ancak o zamandan beri GPU programlama için birincil ortam olarak CuTeDSL'i benimserdik.

Matematiksel bir perspektiften bakıldığında, çıkarım çekirdekleri genellikle karmaşık değildir. Sonuç olarak, bir çekirdek geliştirilirken karmaşıklık genellikle hangi hesaplamanın yapılacağını ifade etmekte değil, donanımın bunu tam olarak nasıl yapması gerektiğinde yatar. En iyi çekirdekleri oluşturmak için, alttaki donanım üzerinde doğru soyutlamalar sunan bir dile ihtiyacımız var. Doğru dil, hesaplamayı gerçekleştirmek için üst düzey işlemleri belirtmemizi, bunları geniş bir yapılandırma Alanındaki tüm noktalar üzerinden agresif bir şekilde özelleştirmemizi ve optimize etmemizi sağlar.

CuTeDSL, CuTe yerleşim cebiri (layout algebra) ve MLIR üzerine kurulmuş NVIDIA GPU'lar için Python tabanlı bir DSL'dir. Düşük seviyeli donanım ilkelleri üzerindeki kontrolünü korurken, son derece optimize edilmiş PTX'e tam zamanında (just-in-time) derleme yapar. Geliştirme hızını engelleyen yüksek derleme süreleri olmadan, CUTLASS C++ ile aynı özellikleri ve performansı sunmayı amaçlar. İncelediğimiz veya birlikte çalıştığımız diller arasından (CUDA, cuTile, TVM IR ve Triton dahil), çekirdeklerimizin geliştirilmesi için en uygun olanı olduğunu düşünüyoruz.

Genel olarak, Triton veya cuTile kadar üst düzey ve kısa olmasa da, CuTe Yerleşim Cebiri ve kompozisyonlarının (cute.Tensor aracılığıyla sunulur) makul ikameler olduğunu gördük. Ek olarak CuTeDSL, tüm uç durumlarda en yüksek performansı elde etmek için gereksinim duyduğumuz kontrol düzeyini sağlayarak donanım ilkellerine kısıtlamasız erişim sunar. CUDA ile, her ek parametre derleme sürelerine makul olmayan bir maliyet veya derleme sistemine karmaşıklık eklediğinden, derleme zamanı uzmanlaşması fırsatları sınırlıydı. Buna karşılık, Tam Zamanında Derlenen (JIT) bir DSL ile yineleme hızından ödün vermeden daha agresif bir şekilde uzmanlaşabiliriz. Ara adım olarak CUDA kullanan TVM IR'ın aksine, PTX'e düşen MLIR tabanlı boru hattının, ara adım olarak CUDA yayan bir boru hattından daha sağlam ve hata ayıklaması daha kolay olduğunu gördük. Ek olarak, PTX talimatlarını sarmalamak için özel işlemlerin (ops) kullanımıyla, alttaki donanımın tüm özelliklerinden tam olarak yararlanabileceğimizi gördük.

CUDA çekirdeklerinin derleme maliyetinin büyük kısmı tipik olarak çekirdeklerin şablon genişletme aşamasında harcanır, çünkü PTX aşamasından itibaren optimizasyon ve derleme birkaç saniye içinde tamamlanabilir. Birçok derleme zamanı parametresi ile derlemenin bu aşaması, önemsiz olmayan bir çekirdek için birkaç dakikaya çıkabilir. Bu darboğaz, C++ şablon işleme sırayla gerçekleşirken, paralel derleme fırsatlarının yalnızca daha sonra, IR ve PTX düzeylerinde ortaya çıkması gerçeğiyle vurgulanır. Şablon örneklemelerini ayrı derleme birimlerine taşıyarak bunu ele alacak püf noktaları mevcut olsa da, bunlar derleme sistemlerinde önemli ince ayarlar ve ağır makroların kullanılmasını gerektirir. Buna karşılık, Tam Zamanında (JIT) derlenen CuTeDSL çekirdekleri karmaşık derleme kuralları gerektirmez ve derleme sürelerini engellemez; çünkü geliştirme döngüsü yapılandırma Alanındaki tek bir noktaya odaklanabilirken, daha sonra dikkatli birim testleri çekirdeğin tüm tasarım Alanını kapsayabilir.

JIT derlemesi maliyetleri geliştirme zamanından çalışma zamanına taşısa da ek yük fazlasıyla kabul edilebilirdir. Bir model dağıtıldığında, yapılandırma Alanının tamamından yalnızca birkaç noktayı etkinleştirerek çıkarım motorunun başlatılması/ısıtılması sırasında saniye mertebesinde ek yükler ortaya çıkarır. Buna karşılık, bir çekirdeğin tüm yapılandırmalarını önceden derlemek birkaç dakika sürebilir. Ek olarak, JIT ile derlenen yapıtlar önbelleğe alınarak ek yükler daha da azaltılabilir.

CUDA'nın yüksek derleme maliyeti, CuTe ve CUTLASS'ta olduğu gibi şablon tabanlı dağıtım ve uzmanlaşma kullanıldığında hataların kökenini izlemenin zorluğuyla daha da artar. Bir DSL ile hatalar onu oluşturan kod satırına yönlendirilebilir, bu da hata ayıklamayı önemli ölçüde kolaylaştırır. Ek olarak, ara MLIR aşamaları aracılığıyla veya üretilen PTX'e eklenen hata ayıklama bilgilerine güvenerek oluşturulan kodu onu oluşturan DSL ifadesiyle eşlemek çok daha basittir. DSL'in bu yönleri sayesinde, çekirdek geliştirirken kendimizi çok daha üretken bulduk.

Çekirdekler modeller arasında birçok benzerliği paylaşsa da, yine de performans üzerinde etkisi olabilecek bir dizi yapılandırma ayarıyla donatılmalıdırlar. İlgili parametreler şunlardır:

Gizli Boyut / Baş Boyutu / Dikkat Başı Sayısı (Hidden Dimension/Head Dimension/Number of Attention Heads)

Giriş/Çıkış veri türleri

Niceleme (quantization) şeması

Derleme zamanı sabitleri (sapmalar, aktivasyon fonksiyonları vb.)

Motoru, modeli, çekirdek kitaplığını ve hizmetleri dahil olmak üzere ROSE mimarisi.

Derleme zamanı uzmanlaşması hem performans hem de hata ayıklama için kullanışlıdır. Özellikle GPU'larda olduğu gibi boru hatlı bir işlemcide, var olmayan bir dal, her zaman alınan veya hiç alınmayan bir dalikinden daha iyidir. Ek olarak, bir çekirdeği belirli bir yapılandırma ayarları seti için derlemek, bir araştırma sırasında cuda-gdb kullanılırken analiz edilmesi ve hata ayıklanması daha kolay olan daha az PTX/SASS talimatı ile sonuçlanır.

Belirli bir toplu iş boyutu ve yapılandırma için, DRAM'e giden beklenen trafik ve aritmetik işlem sayısı tahmin edilebilir. Ardından, optimum ızgara (grid) tahsisi, senkronizasyon ve belirli donanım özelliklerinin kullanımı için duruma göre optimum kararlar alınabilir. CuTeDSL, kabul edilemez bir derleme zamanı ek yüküne katlanmadan derin uzmanlaşma için bileşenleri tanımlamamıza ve yeniden kullanmamıza olanak tanır.

Genel olarak, çıkarım çekirdeklerinin performansı iki ayrı kullanım durumu için optimize edilmelidir:

Ön dolgu (Prefill): küçük toplu iş boyutu (az sayıda parça, diziler), dizi başına çok sayıda belirteç. Hem gömmeler hem de LLM'lerin giriş belirteçleri için iş hacmi maksimize edilmelidir. On ila yüzlerce milisaniye cinsinden gecikme.

Decode: küçükten büyüğe toplu boyutlar, her dizi için bir (tahmine dayalı çözme yok) ila birkaç token (taslak-hedef spekülasyonu, çoklu token tahmini). Çekirdek başlatma gecikmesine yakın gecikme.

GEMM ve dikkat (attention) için, ön dolgu aşaması çok sayıda belirteç üzerinde çalışır ve performans cihazın hesaplama iş hacmi ile sınırlıdır. Kod çözme, küçük toplu iş boyutlarında aynı işlemi içerir ve yürütme süreleri, cihazın bellek bant genişliği tarafından kısıtlanarak, aktivasyonları ve ağırlıkları DRAM'den akış çoklu işlemcilerine (streaming multiprocessors) ve tensör çekirdeklerine yükleme maliyeti tarafından domine edilir. Toplu iş boyutu veya bir toplu işteki toplam belirteç sayısı nispeten net bir gösterge olmasa da, ön dolgu ve kod çözme arasındaki geçiş eşiği daha inceliktir; çünkü bellek bant genişliği aynı zamanda modelin gizli boyutu ve aktivasyonun veri türü tarafından belirlenir. Algoritmik farklılıklar veya ızgara/blok atamaları bunlardan birine avantaj sağlayacaksa, çekirdekler sonuç olarak ön dolgu ve kod çözme uygulamaları olarak bölünür.

Ön dolgu ve kod çözme çekirdekleri arasında evrensel olarak farklı olan parametre, bir bloktaki warp sayısıdır (warp başına 32 thread olan warp boyutuna bölünmüş thread sayısı). Kayıt (register) tahsisine yardımcı olabileceğinden, bu parametreyi bir derleme zamanı sabiti yapmak arzu edilir. CUDA'da bu genellikle __launch_bounds__ aracılığıyla thread sayısını belirterek ve warp sayısını bir şablon parametresi olarak sağlayarak yapılır:

cpp

template <size_t NUM_WARPS>
__global__ void 
__launch_bounds__(NUM_WARPS * WARP_SIZE, 1)
some_kernel(...) {

}

Bir bloktaki tüm thread'lerin erişebileceği register sayısı bir donanım kısıtlamasıdır (Hopper ve Blackwell'de 64K). Thread sayısı önceden biliniyorsa, derleyici her thread'e tam olarak kaç register atayabileceğini bilir ve daha agresif bir şekilde optimize edilmiş çekirdekleri derlemekte başarılı olabilir. Ancak bu, çekirdeğin farklı warp boyutları için yeniden derlenmesi gerekeceğinden, bir derleme zamanı maliyetiyle gelir. CuTeDSL ile bu artık bir sorun değildir, çünkü yapılandırma Alanının tamamı için derleme maliyetini önceden ödememiz gerekmez. DSL ayrıca blok boyutunun sabit bir ifade olmasını bekler ve CUDA'daki __latch_bounds__ tarafından ima edilen optimizasyonu örtük olarak gerçekleştirir.

Warp sayısını sabit yapmanın yanı sıra, bunu ön dolgu ve kod çözme arasında değiştirmek arzu edilir. Bir GPU'nun bellek bant genişliğinden tam olarak yararlanmak için, blok başına her zaman yüksek sayıda warp atamak yeterli değildir; çünkü tam bant genişliğine ancak mevcut Akış Çoklu İşlemcilerinin (SM) olabildiğince çoğu kullanılarak ulaşılabilir. Sonuç olarak, belirli sayıda belirteç ile bilinen veri türleri ve gizli boyutlar için, üretilen kodun 128 bitlik yükleme/depolama talimatlarını (Blackwell'de 256 bit) kullanmasını sağlamak üzere önce ideal bir thread başı vektör boyutu (bfloat16 için 8, float32 için 4) seçeriz. Vektör boyutunu bilerek, çekirdeğin gerektirdiği warp sayısını hesaplayabiliriz. Önce SM sayısını maksimize etmeye çalışırız: SM'lerden daha az warp varsa, blok başına tek bir warp seçilir. Bu genellikle çekirdeğin kod çözme varyantıdır. Yeterli warp varsa, performans açısından artık bir iyileşme görmediğimiz üst sınırı kıyaslayarak warp sayısını kademeli olarak 8, 16 veya 32'ye çıkarırız.

Triton bu optimizasyonların bazılarını kaput altında gerçekleştirirken, üst düzey bir derleyicinin eksiklikleriyle kısıtlanmadan, CuTe yerleşim cebirini dikkatli bir şekilde kullanarak benzer sonuçlar elde edebileceğimizi gördük. Özellikle warp altı yatay indirgemeler söz konusu olduğunda, CuTeDSL'in optimal kod üretmek için gereken ince taneli kontrolü bize sunduğunu gördük.

Bazı durumlarda, işi olabildiğince çok bloğa dağıtmak için, katılan tüm blokları senkronize etmek üzere cooperative_groups::this_grid().sync() kullanılabilir. Kısıtlamalar geçerlidir: bir blok yürütmeye başladıktan sonra bir SM'den zamanlaması kaldırılamayacağından, ileriye dönük ilerlemeyi garanti etmek için blok sayısı cihazda mevcut SM sayısını aşamaz; bu da herhangi bir zamanda tüm blokların senkronizasyonda yer alan küresel bellek konumlarını güncelleyebilmesini ve yoklayabilmesini sağlar. Bu kısıtlama cudaLaunchCooperativeKernel tarafından zorlanır. Bir ızgara bariyerinin maliyeti Hopper'da yaklaşık 3us'dir; buna bazı durumlarda örtük bir bellek bariyeri aracılığıyla önceki tüm bellek işlemlerini temizleme maliyeti de dahildir. Genel olarak bu, bir çekirdeği iki başlatmaya bölmenin ek yüküyle karşılaştırılabilirdir ve performans açısından hassas bir tasarım kararına yol açar. Programmatic Dependent Launch (PDL), Hopper'dan itibaren başlatma ek yüklerini azaltabilmesine rağmen, kod çözme çekirdekleri için ek yükten kaçınmaya değer.

İş yüklerinden biri SM sayısından daha fazla blok gerektiriyorsa, çekirdek başlatmalarının örtük bariyerine güvenilerek çekirdek bölünmelidir. Bu öncelikle ön dolgu için geçerlidir. Tüm iş yükleri için ideal olan blok sayısı ızgaraya sığıyorsa, daha az sayıda belirteç üzerinde çalışan kod çözme çekirdeklerinde olduğu gibi bir ızgara bariyeri tercih edilir.

CuTeDSL'de bu kararlar ana bilgisayar (host) tarafındaki başlatma kodunda verilebilir ve aşağıdaki örnekte gösterildiği gibi, çekirdekler gerektiğinde ızgara bariyerleriyle oluşturulabilir ve araya sokulabilir. Böyle hibrit bir strateji ile kod çözme çekirdekleri minimum gecikmelere ulaşacak şekilde, ön dolgu çekirdekleri ise maksimum iş hacmi için optimize edilebilir. CUDA tabanlı çözümlerin aksine, bu mantık daha kısa bir şekilde ve daha az derleme zamanı ek yüküyle tanımlanabilir.

python

@cute.kernel
def dispatch_kernel_decode(...):
	route_tokens(...)
	grid_barrier(grid_barrier_ptr)
	copy_tokens_to_experts(...)

@cute.kernel
def dispatch_kernel_prefill_route(...):
	route_tokens(...)
	
@cute.kernel
def dispatch_kernel_prefill_copy(...):
	copy_tokens_to_experts(...)
	
@cute.jit
def dispatch(...):
	num_route_blocks = ...
	num_copy_blocks = ...
	
	if num_copy_blocks < NUM_SMS:
		dispatch_kernel_decode(...).launch(
			grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
		)
	else:
		dispach_kernel_prefill_route(...).launch(
			grid=(num_route_blocks, 1, 1)
		)
		dispatch_kernel_prefill_copy(...).launch(
			grid=(num_copy_blocks, 1, 1)
		)

MoE all-to-all çekirdeklerimiz söz konusu olduğunda, aşağıda detaylandırıldığı gibi, ızgara bariyerleri olan veya olmayan çekirdekler oluşturma esnekliğine sahip olmak ve bunları derleme zamanında optimize etmek, CuTeDSL aracılığıyla önemli performans kazanımları elde etmemize yardımcı oldu.

CuTeDSL ile oluşturduğumuz ilk çekirdek, küçük boyutlar için optimize edilmiş özel bir RMS norm çekirdeğiydi. Son Qwen modellerinin öz-dikkat (self-attention) katmanında kullanılan QK normu, 64, 128 veya 256 öğe üzerinde baş başına (per-head) uygulanır. Hem paylaşımlı bellek hem de registerlar aracılığıyla yatay bir indirgeme gerektiren, binleri aşan gizli boyutlarda çalışan normal RMS normunun aksine, QK normu yalnızca warp'lar arasında indirgeme gerektirir ve birden fazla başlığı tek bir warp'a sıkıştırma olasılığı bulunur. Bu, tüm baş boyutları ve veri türleri için en iyi belirteç-warp atamalarını bulmak amacıyla derleme zamanında agresif bir şekilde optimize edebileceğimiz bir kullanım durumu haline geldi.

RMS normu oldukça anlaşılır bir işlemdir, ancak bir vektörün karesi alınmış normunu hesaplamak için yatay bir toplam gerektirir:

cpp

r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

Dikkat başlıkları tipik olarak küçük olduğundan ve thread başı işlemler büyük ölçüde vektörleştirildiğinden, tek bir warp kare normunu hesaplayabilir. Bellek bant genişliğini maksimize etmek için her thread, 8 bfloat16 değerine karşılık gelen 128 bitlik girişi işlemelidir. Sonuç olarak, bir başlığı işlemek için yarım warp (half-warp) yeterlidir ve kare norm, shuffle_sync_bfly tarafından gerçekleştirilen yatay bir indirgeme kullanılarak belirlenebilir. Baş boyutundan ve veri türünden başlayarak, warp başına baş sayısını belirleyebilir ve derleme sırasında argümanları ve indirgeme turu sayısını belirtebiliriz.

Ardından, otomatik vektörleştirilmiş yükleme/depolamalara güvenerek ve warp indirgemesi için alt düzey mimariye özel işlemler kullanarak nispeten az sayıda kod satırıyla özel bir çekirdek oluşturabiliriz. Çekirdek üzerinde çalışılacak başlığı (Q veya K) belirler ve uygulanacak ilgili ağırlıkları tanımlar. Baştan tüm bir vektör yükler ve önce yerel olarak toplar. Ardından, baş başına thread sayısına bağlı olarak kare norm hesaplanır. Son olarak, otomatik olarak vektörleştirilmiş RMS normu belirlenir ve belleğe geri yazılır. WEIGHT_BIAS ve EPS gibi modele özel parametreler sabit katlanmıştır (constant-folded).

python

@cute.kernel
def _qk_norm_kernel(
    q: cute.Tensor,
    k: cute.Tensor,
    q_weight: cute.Tensor,
    k_weight: cute.Tensor,
    WEIGHT_BIAS: cutlass.Constexpr[float],
    EPS: cutlass.Constexpr[float],
) -> None:
    # Find the block and thread indices.
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    tidx, _, _ = cute.arch.thread_idx()

    _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
    _, num_k_heads, _ = k.shape
    head_dim = threads_per_head * vecs_per_thread * vec_size

    # Determine the thread and the vector to work on.
    vec_id = (bidx * bdim + tidx) // threads_per_head
    thread_id = tidx % threads_per_head

    # Determine the token and head indices.
    token = vec_id // (num_q_heads + num_k_heads)
    head_offset = vec_id % (num_q_heads + num_k_heads)
    is_q = head_offset < num_q_heads

    # Determine which tensor to work with.
    thread_coord = (thread_id, None, None)
    head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)

    # Load the head into registers.
    x_ref = q[head_coord] if is_q else k[head_coord]
    w = (q_weight if is_q else k_weight)[thread_coord].load()

    # Compute the squared sum locally.
    x = x_ref.load()
    x_f32 = x.to(cute.Float32)
    x_square = x_f32 * x_f32
    x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)

    # Sum up across the warp.
    mask = activemask()
    start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
    for offset in [start >> i for i in range(5) if (start >> i) > 0]:
        x_square_sum += cute.arch.shuffle_sync_bfly(
            x_square_sum,
            offset=offset,
            mask=mask,
        )

    # RMS norm.
    rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
    x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
    x_ref.store(x_norm)

Normal bir RMS norm çekirdeği ile karşılaştırıldığında, bu özel varyant hem ön dolgu hem de kod çözme için 2-3 kat daha iyi iş hacmi ve gecikme sunar.

MoE dağıtımı ve birleştirmesi, giriş belirteçlerini yönlendirildikleri uzmanlara karıştırarak, MoE GEMM çekirdekleri ilgili matris çarpımlarını gerçekleştirdikten sonra ağırlıklı bir biriktirme gerçekleştirir. Çeşitli ürünler çeşitli boyutlardaki MoE modellerine güvendiğinden, performans özellikle ilgi çekicidir. Dağıtım türlerine dayanarak, optimize edilecek farklı çekirdekleri ayırt ederiz:

TP=1: Tek bir GPU için yerel çekirdek. Dağıtım, belirteçleri ilgili uzmanlara karıştırır ve birleştirme bunları biriktirir. Öncelikle H100 üzerindeki Qwen3 30B dağıtımları için kullanılır.

DP=1, TP ≠ 1: Tensör paralelliğine sahip, tek düğümlü çekirdek. Giriş çoğaltılır, dağıtım mevcut derece (rank) için belirteçleri seçer ve birleştirme bunları tüm dereceler genelinde biriktirir. NVLink veya NVLS üzerinden biriktirme. Büyük ön doldurucular için kullanılır.

DP ≠ 1, NVLink: İsteğe bağlı tensör paralelliğine sahip veri paralelli (data-parallel) çekirdek. Her derece, belirli uzmanları barındıran diğer her dereceye belirteçler gönderir. Yönlendirme bilgileri toplulaştırılır, belirteçler NVIDIA NVLink üzerinden aktarılır. Veri paralelli H200 veya GB200 dağıtımları için kullanılır.

DP ≠ 1, InfiniBand: Hem NVLink hem de NVIDIA Quantum InfiniBand üzerinden çalışan Eşler Arası (Peer-to-Peer) veri paralelli ve tensör paralelli çekirdek. Çok düğümlü Hopper dağıtımları için kullanılır.

Yerel çekirdek, yönlendirme tablosuna dayanarak her uzmana atanan belirteç sayısını sayar ve her uzmana ait belirteç listesi içindeki her belirteçe benzersiz bir indeks atar. Yönlendirme adımından sonra belirteçler, bir MoE GEMM çekirdeğine iletilmek üzere isteğe bağlı olarak doldurulmuş (padded) tamponlara (buffers) kopyalanır. Belirteç konumlarını kullanarak birleştirme çekirdeği daha sonra register'larda bir biriktirme gerçekleştirir ve son belirteçleri bir çıkış tamponuna yazar.

Dağıtım aşaması ön dolgu ve kod çözme arasında değişiklik gösterir: az sayıda belirteç için okunacak çok az yönlendirme bilgisi vardır, bu nedenle bunu tek bir blokta işler ve ızgaranın geri kalanını sinyalleriz, bireysel bloklar belirteç öbeklerini kopyalar. Belirteç ataması, paylaşımlı bellekteki sayaçların atomik olarak artırılmasıyla belirlenir. Ön dolgu daha fazla yönlendirme bilgisini işlemesi gerektiğinden, yönlendirme ayrı bir çekirdekte birden fazla bloğa bölünür. İndeksler önce paylaşımlı bellekteki sayaçlar kullanılarak yerel olarak atanır ve ardından her belirtecin konumunu belirlemek için küresel bellekteki tüm bloklar genelinde toplulaştırılır. Bu çekirdeğin ızgarası, ızgara tabanlı bariyerleri kullanabilmek için SM sayısına kelepçelenir (clamped). Sonraki bir çekirdek, mevcut bellek bant genişliğinden tam olarak yararlanmak için çok daha büyük sayıda blok ve warp arasında başlatılarak kopyalama işlemlerini gerçekleştirir.

Birleştirme çekirdeği, yönlendirme bilgilerine ve dağıtım tarafından hesaplanan uzaklıklara (offsets) güvenerek vektörleştirilmiş biriktirme gerçekleştirdiği için nispeten basittir. Hesaplama, tüm veri türleri ve gizli boyutlar için verimli koda derlenen basit CuTe tensörleri aracılığıyla ifade edilebildiğinden, CuTeDSL burada değerli olduğunu kanıtladı.

Tensör paralelliğine sahip bir dağıtımda, yönlendirme bilgileri ve giriş belirteçleri çoğaltılır. Yerel dağıtım çekirdeği, yalnızca mevcut derecede bulunan uzmanlara atanan belirteçleri yönlendirmek ve seçmek için ufak bir değişiklikle tamamen yeniden kullanılabilir. Bazı parametreler ve derleme zamanı uzmanlaşmasıyla, yerel dağıtım çekirdekleri yeniden kullanılır.

Derece sayısı genellikle bir belirtecin yönlendirildiği uzman sayısından daha düşük olduğundan, ortalama durumlarda her belirtecin her dereceye yönlendirilen kopyalarından en az birine sahip olduğu varsayılabilir. Bu durumda yerel birleştirme çekirdeği, bir belirtecin yerel kopyalarını toplulaştırmak için yeniden kullanılabilir; bunu katılan dereceler arasında toplama yapmak için bir all-reduce takip eder. Ön dolgu çekirdeği durumunda ayrı çekirdekler başlatırız, ancak kod çözme her iki işlemi tek bir başlatmada birleştirir.

All-reduce için, yalnızca iki rank katılıyorsa düz NVLink kullanırız; 4 veya daha fazla rank devreye girdiğinde ise verimde %10-20 artış için NVLS’ye geçeriz. DSL sayesinde, MoE ile sıradan all-reduce çekirdekleri arasındaki indirgeme mantığının ve peer buffer değişimlerinin büyük bölümünü paylaşabiliriz. Bu, donanım üzerinde yakın kontrolün yararlı olduğu bir durumdur; çünkü NVLS için multimem.ld_reduce ve multimem.st üzerinde DSL işlemleri gerekir. İndirgeme işleminin uygulaması büyük ölçüde NCCL testlerindeki all_reduce örneğini izler.

Veri Paralelli çekirdekler, birden fazla cihaz için belirteçleri toplulaştırmak amacıyla NVLink'e güvenir. Dağıtım aşamasında, eşler belirteçleri değiştirmeden önce, gönderdikleri her belirteç için her derecedeki hedef uzaklığı belirlemek üzere önce yerel olarak yönlendirme yapar. Ardından, çakışmaları önlemek için her gönderen için ayrılan alma tamponlarına (receive buffers) belirteçleri kopyalamaya başlarlar. Tamamlandığını sinyallemek için belirteç başı (per-token) release-store yazmaları kullanılır. Alıcı tarafında, tüm belirteçleri MoE GEMM için tek bir tamponda yoğunlaştırabilmek amacıyla yönlendirme bilgileri toplulaştırılır. Alıcılar, yazmaların tamamlanmasını beklemek için acquire-load ile eş tarafından yazılan bayrağı yoklar.

Birleştirme tarafında, belirteçler NVLink üzerinden kaynak derecelere kopyalanır ve alımı belirten bir belirteç başı bayrak bulunur. Birleştirme çekirdeği, adreslemede hafif farklılıklar ve bazı ek senkronizasyon mantığı ile yerel ve tek düğümlü varyantlara benzer.

Bu çekirdeği uygulamak için hem sistem hem de cihaz düzeyinde kapsamlandırılmış release-acquire yüklemelerini ve saklamalarını (stores) sunan bir PTX sarmalayıcıları kitaplığı oluşturduk. Senkronizasyon bayraklarının ayrıntı düzeyini artırarak ve NVLink üzerinden belirteç göndermek ve almak için daha fazla blok ayırarak, önceki CUDA çekirdeklerimize göre 2 kat iyileşme elde ettik.

Eşler Arası çekirdeklerimiz daha önce ConnectX-7 InfiniBand bağdaştırıcılarıyla eşleştirilmiş Hopper GPU'larda son teknoloji performans elde etmişti. Performansta daha fazla kazanç elde etmek için o zamandan beri birleştirme çekirdeklerini CuTeDSL'e taşıdık. CUDA sürümleri halihazırda büyük ölçüde derleme zamanı uzmanlaşmasına güvendiğinden, depomuzdaki derleme zamanı darboğazları olduklarını kanıtladılar ve bazı parametreleri atık bırakmaya zorladılar. JIT derlemeli DSL'e geçerek, daha fazla parametre için uzmanlaştık ve ızgara tahsislerine ince ayar yaptık; bu da gecikmede yaklaşık %3'lük bir artış için ön dolgu ve kod çözme arasında ek farklar getirdi.

ROSE çekirdekleri perspektifinden Blackwell desteği iki kategoriye ayrılır: mimariden bağımsız (architecture-agnostic) ve mimariye özel (architecture-specific) çekirdekler. Tensör çekirdeklerine güvenen mimariye özel çekirdekler uyum sağlamak için açık çalışma gerektirir. Ancak DSL'in kullanımı sayesinde mimariden bağımsız çekirdekler önemsiz bir şekilde uyum sağlar.

Bir Blackwell GPU, Hopper'ın 132'sine karşılık 160'a kadar SM'ye sahiptir. Çoğu çekirdek için bu bir derleme zamanı sabitidir ve ana bilgisayar tarafındaki başlatıcı kodundaki mantık, çalıştırılacak optimum (ön dolgu veya kod çözme) çekirdeği belirlemek için bunu hesaba katar. Ek olarak Blackwell, bellek bant genişliği kullanımını iyileştirmek için vektör genişliği (CuTe tensörlerinin son boyutu) artırılarak kolayca yayılabilecek 256 bitlik yükleme talimatları da sunar.

Blackwell, önemli sayıda niceme şeması için donanım desteği sunar (MXFP8, MXFP4, NVFP4). Bunlar için destek öncelikle FP8 blok bazlı ölçeklendirme için mevcut çekirdeklere ince ayar yapılarak uygulanmıştır. Yeni modlar, derleme zamanı yapılandırma Alanını patlatmadan desteklenebileceğinden JIT desteği burada çok önemlidir. Geçmişte Triton'un birkaç öğe üzerinden, bir warp'ın yarısı veya çeyreği üzerinde yatay indirgemeleri ifade etmekte tuhaf olduğunu gördük. NVFP4 ve MXFP8/MXFP4 için sırasıyla 16 veya 32 öğe üzerinden mutlak bir maksimum değer hesaplanırken optimal performans elde etmek için indirgemeler üzerinde daha ince taneli kontrol çok önemliydi.

Nispeten yeni bir dil olarak LLM'ler, ön eğitimleri sırasında CuTeDSL koduna çok az maruz kalmıştı. Ek olarak, aktivasyonlar ve normlar için optimize edilmiş çekirdeklere sahip mevcut kodumuzun stili, TMA ve mimariye özel matris çarpımı talimatları için yoğun bir şekilde optimize edilmiş mevcut CUTLASS veya Quack örneklerinden biraz farklıydı. LLM'lerin çekirdek yazmada yardımcı olabileceği bir noktaya ulaşmak kademeli bir süreçti.

Ajanlara çekirdek dönüştürme konusunda güvenmeden önce, tutarlı kalıpları izleyerek manuel olarak ilk bir çekirdek ve yardımcı fonksiyonlar kümesi oluşturduk. Kürop ilerledikçe, önemli bir manuel yeniden çalışmayı takip ederek CUDA'dan CuTeDSL'e kaba çeviriler gerçekleştirebildik. Özellikle modellerin işaretçi-tensör (pointer-to-tensor) dönüşümleri ve vektörleştirme ile mücadele ettiğini, bir testin derlenmesini sağlamayı hedeflerken tür hatalarını düzeltmeye çalışarak takılıp kaldığını gördük. Çekirdeklerin ızgara tabanlı doğasını çok az anladıklarını göstererek thread ve ızgara senkronizasyonunu karıştırdılar. En sinir bozucu olanı, kendi CuTeDSL alternatiflerimize karşı derlemek yerine CUDA'ya özel kitaplıkların var olmayan çevirilerini halüsinasyon görme eğilimi olmasıydı. Zamanla, CuTeDSL çekirdeklerinin kürepesi büyüdükçe LLM'lerin iskelet kodunun çoğunu oluşturabileceğini ve boşlukları doldurabileceğini, böylece üst düzey tasarıma odaklanmamıza olanak tanıdığını gördük.

CUDA ile karşılaştırıldığında CuTeDSL'in çok daha iyi derleme zamanı performansı sunduğunu, daha fazla parametre için uzmanlaşmamıza ve performans iyileştirmeleri elde etmemize olanak tanıdığını görüyoruz. Triton'a kıyasla belirli soyutlamaların eksikliği nedeniyle çekirdekler biraz daha ayrıntılı olsa da, donanım üzerinde daha ince taneli kontrol daha hızlı çekirdekler oluşturmamıza izin verdiği için karmaşıklığın haklı olduğunu görüyoruz.

Yazıldığı sırada, CUTLASS C++ ve CuTeDSL çekirdekleri arasındaki performans boşluğu minimum düzeydedir ve yüksek performanslı GEMM çekirdeklerini cihazlar arası iletişim, normalleştirme veya aktivasyonlarla birleştirmek ve harmanlamak için heyecan verici fırsatlar açmaktadır. Bu DSL'e bağlı kalarak, gömmelerden LLM'lere kadar iş yüklerimiz genelinde iş hacmini iyileştirmek için cihazlarımızdan daha fazla performans çıkarabileceğiz.

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

BibTeX

@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}

Perplexity AI'da Çıkarım

Neden CuTeDSL?

Geliştirici deneyimi

Geniş yapılandırma Alanı

Ön dolgu ve kod çözme uzmanlığı

Izgara (grid) senkronizasyonu

CuTeDSL Çekirdekleri

QK Normu

Mixture-of-Experts Dağıtımı/Birleştirmesi (Dispatch/Combine)

Yerel (Local)

Tek düğümlü (Singlenode)

Veri Paralelli (Data Parallel)

Eşler Arası (Peer-to-Peer)

Blackwell Desteği

Yapay Zeka Destekli Çekirdek Oluşturma

Sonuç ve Gelecek Çalışmalar

Referanslar

Atıf