CuTeDSL di Perplexity

Di Perplexity, kami menggunakan Runtime-Optimized Serving Engine (ROSE) buatan internal untuk melayani model mulai dari penyematan hingga LLM dengan satu triliun parameter. Berada di balik semua API, termasuk Sonar, Search, dan Embeddings, runtime kami dibangun untuk memberdayakan para peneliti guna mengeksplorasi model dan arsitektur yang paling sesuai, sekaligus dengan cepat menerapkan kandidat terbaik ke dalam produksi. Untuk mencapai hal ini, kita perlu mampu menerapkan, menggelar, dan mengoptimalkan inferensi untuk arsitektur mutakhir dengan cepat. Dalam postingan blog ini, kami mengeksplorasi bagaimana kernel CuTeDSL yang diintegrasikan ke dalam mesin inferensi kami memungkinkan kami membangun kernel GPU yang diperlukan untuk mendongkrak performa model ke tingkat puncak pada GPU NVIDIA.

Model kustom, yang dibangun oleh berbagai tim di Perplexity, di-host secara internal pada GPU NVIDIA Hopper dan Blackwell menggunakan mesin inferensi ROSE yang dibangun oleh tim AI Inference. ROSE awalnya dibangun untuk melayani model Llama yang dikustomisasi dengan lebih baik untuk dekode bahasa dan klasifikasi, dengan antarmuka yang kompatibel dengan NVIDIA Triton Inference Server. Seiring waktu, ROSE berkembang menjadi mesin inferensi berfitur lengkap yang melayani LLM dalam skala besar, serta model perankingan, klasifikasi, penskoran, dan model berbasis transformer lainnya untuk berbagai tim guna mendukung produk-produk Perplexity. Tujuan desain ROSE adalah agar tetap sederhana: seiring waktu, upaya diinvestasikan untuk mencapai performa puncak, tetapi prioritas utamanya adalah beradaptasi dengan arsitektur dan kasus penggunaan model kami yang terus berkembang.

Ruang konfigurasi kernel di seluruh tipe data, dimensi head, konstanta waktu kompilasi, dan pengaturan peluncuran.

Secara tingkat tinggi, ROSE terdiri dari engine yang mengadaptasi model ke antarmuka tempat klien dapat berinteraksi. Mesin-mesin ini sangat dapat digunakan kembali, menangani penjadwalan dan pengelompokan permintaan, serta mengambil token terdekode, penyematan, atau skor dari model yang mendasarinya. Inisialisasi perangkat, saluran komunikasi antarproses dan antarnode, serta pemuatan bobot semuanya ditangani olehnya. Untuk LLM, mesin ini bertanggung jawab atas pengelompokan, pemotongan, dan pengambilalihan sampel (sampling). Mesin ini juga bertanggung jawab atas alokasi penyimpanan KV, untuk perhatian penuh (full attention) maupun linier, serta pencocokan prefiks. Mesin penyematan (embedding engines) lebih sederhana karena hanya melakukan pengelompokan secara online. Eksekusi dioptimalkan sedemikian rupa sehingga mesin menyiapkan putaran pekerjaan berikutnya sementara model berjalan pada akselerator secara paralel.

ROSE mengekspos banyak koleksi lapisan kustom untuk pemodelan, membungkus kernel GPU yang dibangun dengan CuTeDSL, Triton, CUDA, CUTLASS, dan cuBLAS. Untuk implementasi perkalian matriks dan perhatian (attention), kami terutama mengandalkan kernel NVIDIA dari CUTLASS dan cuBLAS. Sisa alur kerja, termasuk kernel penyematan, norm, MoE, dan aktivasi, merupakan bagian dari pustaka kernel ROSE. Untuk mendukung perangkat keras tempat kami menjalankan inferensi dengan sebaik-baiknya, ROSE dilengkapi dengan primitif komunikasi antarperangkatnya sendiri untuk perutean, pengiriman (dispatch), dan penggabungan MoE. Meskipun kernel lainnya relatif sederhana, kernel tersebut harus dikhususkan secara mendalam untuk mencapai performa puncak pada semua dimensi tersembunyi dan perangkat yang berbeda tempat kernel tersebut berjalan. Selain itu, kami mengandalkan spesialisasi waktu kompilasi untuk mengakomodasi semua tombol dan penyesuaian dari berbagai model (bias bobot vs. tanpa bias bobot untuk norm RMS, fungsi aktivasi, dll.). Akibatnya, kami memutuskan untuk menggunakan DSL guna mengimplementasikannya: awalnya kami membangun ROSE menggunakan Triton. Namun, sejak saat itu kami telah mengadopsi CuTeDSL sebagai lingkungan utama untuk pemrograman GPU.

Dari perspektif matematis, kernel inferensi biasanya tidak rumit. Akibatnya, saat mengembangkan sebuah kernel, kompleksitas umumnya tidak terletak pada pengungkapan komputasi apa yang harus dilakukan, melainkan pada bagaimana tepatnya perangkat keras harus melakukannya. Untuk membangun kernel terbaik, kami memerlukan bahasa yang menawarkan abstraksi yang tepat atas perangkat keras yang mendasarinya. Bahasa yang tepat memungkinkan kami menentukan operasi tingkat tinggi untuk melakukan kompilasi, serta mengkhususkan dan mengoptimalkannya secara agresif di semua titik dalam ruang konfigurasi yang luas.

CuTeDSL adalah DSL berbasis Python untuk GPU NVIDIA yang dibangun di atas aljabar tata letak CuTe dan MLIR. Bahasa ini dikompilasi secara just-in-time menjadi PTX yang sangat dioptimalkan, sekaligus mempertahankan kontrol atas primitif perangkat keras tingkat rendah. CuTeDSL bertujuan untuk menawarkan fitur dan performa yang sama seperti CUTLASS C++, tanpa waktu kompilasi yang mahal yang menghambat kecepatan pengembangan. Dari sekian bahasa yang telah kami survei atau gunakan (termasuk CUDA, cuTile, TVM IR, dan Triton), kami menganggapnya sebagai yang paling sesuai untuk pengembangan kernel kami.

Secara keseluruhan, meskipun tidak setingkat dan sesingkat Triton atau cuTile, kami menemukan bahwa Aljabar Tata Letak CuTe (CuTe Layout Algebra) dan komposisinya (yang diekspos melalui cute.Tensor) adalah pengganti yang wajar. Selain itu, CuTeDSL menawarkan akses tak terbatas ke primitif perangkat keras, memberi kami tingkat kontrol yang kami butuhkan untuk mencapai performa puncak di semua kasus tepi (edge cases). Dengan CUDA, peluang untuk spesialisasi waktu kompilasi menjadi terbatas, karena setiap parameter tambahan menambah biaya yang tidak wajar pada waktu kompilasi atau kompleksitas pada sistem build. Sebaliknya, dengan DSL yang Dikompilasi Secara Just-In-Time (Just-In-Time Compiled DSL), kita dapat mengkhususkan secara lebih agresif tanpa mengorbankan kecepatan iterasi. Tidak seperti TVM IR yang menggunakan CUDA sebagai langkah perantara, kami mendapati bahwa alur kerja berbasis MLIR yang diturunkan ke PTX jauh lebih tangguh dan lebih mudah di-debug daripada alur kerja yang memancarkan CUDA sebagai langkah perantara. Selain itu, dengan menggunakan op kustom untuk membungkus instruksi PTX, kami mendapati bahwa kami dapat memanfaatkan semua fitur perangkat keras yang mendasarinya secara penuh.

Sebagian besar biaya kompilasi kernel CUDA biasanya dihabiskan pada fase perluasan templat (template expansion phase) kernel, karena pengoptimalan dan kompilasi dari tahap PTX dan seterusnya dapat diselesaikan dalam hitungan detik. Dengan banyaknya parameter waktu kompilasi, tahap kompilasi ini dapat membengkak hingga beberapa menit untuk kernel yang tidak sederhana. Hambatan ini semakin diperparah oleh fakta bahwa peluang untuk kompilasi paralel baru muncul di kemudian hari, yaitu pada tingkat IR dan PTX, sementara penanganan templat C++ terjadi secara berurutan. Meskipun ada trik untuk mengatasi hal ini, dengan memindahkan instansiasi templat ke unit kompilasi terpisah, trik tersebut memerlukan penyesuaian signifikan pada sistem build dan penggunaan makro yang rumit. Sebaliknya, kernel CuTeDSL yang dikompilasi secara Just-In-Time (JIT) tidak memerlukan aturan build yang rumit dan tidak menghambat waktu kompilasi, karena siklus pengembangan dapat difokuskan pada satu titik di ruang konfigurasi, sementara pengujian unit yang cermat di kemudian hari dapat mencakup seluruh ruang desain kernel.

Meskipun kompilasi JIT memindahkan biaya dari waktu pengembangan ke waktu proses (runtime), overhead tersebut sangat dapat diterima. Saat sebuah model diterapkan, model ini hanya mengaktifkan beberapa titik dari seluruh ruang konfigurasi, sehingga menimbulkan overhead berskala detik selama permulaan/pemanasan (startup/warmup) mesin inferensi. Sebaliknya, prakompilasi semua konfigurasi dari sebuah kernel dapat memakan waktu hingga beberapa menit. Selain itu, overhead dapat dikurangi lebih lanjut dengan melakukan cacing (caching) pada artefak hasil kompilasi JIT.

Biaya kompilasi CUDA yang besar semakin diperparah oleh sulitnya melacak asal-usul kesalahan saat pengiriman (dispatch) dan spesialisasi berbasis templat digunakan, seperti pada CuTe dan CUTLASS. Dengan DSL, kesalahan dapat diarahkan ke baris kode yang menghasilkannya, sehingga proses penelusuran kesalahan (debugging) menjadi jauh lebih mudah. Selain itu, jauh lebih mudah untuk memetakan kode yang dihasilkan ke pernyataan DSL yang memicunya melalui tahap MLIR perantara atau dengan mengandalkan informasi penelusuran kesalahan yang dilampirkan pada PTX yang dihasilkan. Berkat aspek-aspek DSL ini, kami mendapati diri kami jauh lebih produktif saat mengembangkan kernel.

Meskipun kernel memiliki banyak kesamaan di berbagai model, kernel tersebut harus tetap dilengkapi dengan berbagai tombol konfigurasi yang dapat berdampak pada performa. Parameter yang relevan meliputi:

Dimensi Tersembunyi (Hidden Dimension) / Dimensi Kepala (Head Dimension) / Jumlah Kepala Perhatian (Number of Attention Heads)

Tipe data Masukan/Keluaran (Input/Output)

Skema kuantisasi

Konstanta waktu kompilasi (bias, fungsi aktivasi, dll.)

Arsitektur ROSE, termasuk engine, model, pustaka kernel, dan layanan-layanannya.

Spesialisasi waktu kompilasi berguna untuk performa maupun penelusuran kesalahan. Khususnya dalam prosesor berpipa (pipelined processor), seperti pada GPU, cabang yang tidak ada (nonexistent branch) jauh lebih baik daripada cabang yang selalu diambil (always-taken) atau tidak pernah diambil (never-taken). Selain itu, mengompilasi kernel untuk seperangkat tombol konfigurasi tertentu menghasilkan lebih sedikit instruksi PTX/SASS, yang lebih mudah dianalisis dan di-debug saat menggunakan cuda-gdb selama investigasi.

Untuk ukuran batch dan konfigurasi tertentu, lalu lintas yang diharapkan ke DRAM dan jumlah operasi aritmatika dapat diperkirakan. Selanjutnya, keputusan optimal dapat dibuat berdasarkan kasus per kasus untuk alokasi kisi (grid allocation), sinkronisasi, dan penggunaan fitur perangkat keras tertentu. CuTeDSL memungkinkan kami menentukan dan menggunakan kembali komponen untuk spesialisasi mendalam tanpa menimbulkan overhead waktu kompilasi yang tidak dapat diterima.

Secara umum, performa kernel inferensi harus dioptimalkan untuk dua kasus penggunaan yang berbeda:

Prefill: ukuran batch kecil (beberapa potongan, sekuens), banyak token per sekuens. Throughput harus dimaksimalkan untuk penyematan dan token masukan dari LLM. Latensi berada dalam kisaran puluhan hingga ratusan milidetik.

Decode: ukuran batch kecil hingga besar, satu (tanpa dekode spekulatif) hingga beberapa token per sekuens (spekulasi target draf, prediksi multi-token). Latensi mendekati latensi peluncuran kernel.

Untuk GEMM dan perhatian (attention), tahap prefill berjalan pada sejumlah besar token dan performanya dibatasi oleh throughput komputasi perangkat. Dekode melibatkan operasi yang sama pada ukuran batch kecil, dengan waktu eksekusi yang didominasi oleh biaya pemuatan aktivasi dan bobot dari DRAM ke multiprosesor aliran (streaming multiprocessors) dan inti tensor (tensor cores), yang sebaliknya dibatasi oleh bandwidth memori perangkat. Meskipun ukuran batch atau jumlah total token dalam batch merupakan indikator yang relatif jelas, ambang batas pengalihan (switchover threshold) antara prefill dan decode menjadi lebih halus karena bandwidth memori juga ditentukan oleh dimensi tersembunyi model dan tipe data aktivasi. Akibatnya, kernel dibagi menjadi implementasi prefill dan decode jika perbedaan algoritmik atau penetapan kisi/blok memberikan keuntungan bagi salah satunya.

Parameter yang secara universal berbeda antara kernel prefill dan decode adalah jumlah warp dalam satu blok (jumlah utas, dibagi dengan ukuran warp yaitu 32 utas per warp). Sangat disarankan untuk menjadikan parameter ini sebagai konstanta waktu kompilasi, karena dapat membantu dalam alokasi register. Dalam CUDA, ini biasanya dilakukan dengan menentukan jumlah utas melalui __launch_bounds__ dan menyediakan jumlah warp sebagai parameter templat:

cpp

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

}

Jumlah register yang tersedia untuk semua utas dalam satu blok adalah batasan perangkat keras, yaitu 64K pada Hopper dan Blackwell. Jika jumlah utas diketahui sebelumnya, kompiler tahu persis berapa banyak register yang dapat ditetapkan ke setiap utas dan dapat berhasil mengompilasi kernel yang dioptimalkan secara lebih agresif. Namun, ini harus dibayar dengan biaya waktu kompilasi, karena sekarang kernel harus dikompilasi ulang untuk ukuran warp yang berbeda. Dengan CuTeDSL, ini tidak lagi menjadi masalah, karena kita tidak perlu membayar biaya kompilasi untuk seluruh ruang konfigurasi di awal. DSL ini juga mengharapkan ukuran blok menjadi ekspresi konstan, yang secara implisit melakukan pengoptimalan yang tersirat oleh __latch_bounds__ di CUDA.

Selain menjadikan jumlah warp sebagai konstanta, sangat disarankan untuk memvariasikannya antara prefill dan decode. Untuk memanfaatkan bandwidth memori GPU secara penuh, tidaklah cukup jika hanya selalu menetapkan jumlah warp yang tinggi per blok, karena bandwidth penuh hanya dapat dicapai dengan menggunakan sebanyak mungkin Streaming Multiprocessors (SM) yang tersedia. Oleh karena itu, untuk jumlah token tertentu serta tipe data dan dimensi tersembunyi yang diketahui, pertama-tama kami memilih dimensi vektor per utas yang ideal (8 untuk bfloat16, 4 untuk float32) guna memastikan bahwa kode yang dihasilkan menggunakan instruksi muat/simpan 128-bit (256-bit pada Blackwell). Dengan mengetahui ukuran vektor, kita dapat menghitung jumlah warp yang diperlukan oleh kernel. Pertama-tama, kami mencoba memaksimalkan jumlah SM: jika jumlah warp lebih sedikit daripada SM, satu warp per blok akan dipilih. Ini biasanya merupakan varian decode dari kernel. Jika warp yang tersedia sudah mencukupi, kami secara bertahap meningkatkan jumlah warp menjadi 8, 16, atau 32, serta melakukan pembandingan (benchmarking) untuk menentukan batas atas di mana kami biasanya tidak lagi melihat peningkatan performa.

Meskipun Triton melakukan beberapa pengoptimalan ini di latar belakang, kami mendapati bahwa kami dapat mencapai hasil yang serupa dengan menggunakan aljabar tata letak CuTe secara cermat, tanpa terbatasi oleh kekurangan dari kompiler tingkat tinggi. Terutama ketika reduksi horizontal sub-warp terlibat, kami mendapati bahwa CuTeDSL menawarkan kontrol mendetail (fine-grained control) yang diperlukan untuk menghasilkan kode yang optimal.

Dalam beberapa situasi, untuk mendistribusikan pekerjaan ke sebanyak mungkin blok, cooperative_groups::this_grid().sync() dapat digunakan untuk menyinkronkan semua blok yang berpartisipasi. Ada batasan yang berlaku: karena sebuah blok tidak dapat dijadwalkan ulang (de-scheduled) dari SM setelah eksekusinya dimulai, untuk menjamin kemajuan ke depan, jumlah blok tidak boleh melebihi jumlah SM yang tersedia pada perangkat, guna memastikan bahwa pada titik mana pun semua blok dapat memperbarui dan memindai (poll) lokasi memori global yang terlibat dalam sinkronisasi. Batasan ini diberlakukan oleh cudaLaunchCooperativeKernel. Biaya penghalang kisi (grid barrier) adalah sekitar 3us pada Hopper, ditambah biaya pembersihan (flushing) semua operasi memori sebelumnya melalui penghalang memori implisit dalam beberapa situasi. Secara keseluruhan, ini sebanding dengan overhead pemecahan kernel menjadi dua peluncuran, yang mengarah pada keputusan desain yang sensitif terhadap performa. Meskipun Programmatic Dependent Launch (PDL) dapat mengurangi overhead peluncuran mulai dari generasi Hopper dan seterusnya, untuk kernel decode, menghindari overhead tersebut sangatlah layak dilakukan.

Jika salah satu beban kerja memerlukan lebih banyak blok daripada jumlah SM, kernel harus dipecah, dengan mengandalkan penghalang implisit dari peluncuran kernel. Hal ini terutama berlaku untuk prefill. Jika jumlah blok yang ideal untuk semua beban kerja pas di dalam kisi, penghalang kisi lebih disukai, seperti halnya kernel decode yang beroperasi pada jumlah token yang lebih kecil.

Dalam CuTeDSL, keputusan ini dapat dibuat dalam kode peluncuran di sisi host dan kernel dapat disusun serta diselingi dengan penghalang kisi seperlunya, seperti yang diilustrasikan oleh contoh di bawah ini. Dengan strategi hibrida semacam itu, kernel decode dapat dioptimalkan untuk mencapai latensi minimal dan kernel prefill dapat dioptimalkan untuk throughput maksimal. Tidak seperti solusi berbasis CUDA, logika ini dapat didefinisikan secara lebih ringkas dan dengan overhead waktu kompilasi yang lebih sedikit.

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)
		)

Dalam kasus kernel MoE all-to-all kami, sebagaimana dirinci di bawah ini, memiliki fleksibilitas untuk membangun kernel dengan atau tanpa penghalang kisi dan mengoptimalkannya pada waktu kompilasi membantu kami mencapai peningkatan performa yang substansial melalui CuTeDSL.

Kernel pertama yang kami bangun di CuTeDSL adalah kernel norm RMS khusus, yang dioptimalkan untuk dimensi kecil. QK norm, yang digunakan dalam lapisan self-attention pada model Qwen terbaru, diterapkan secara per-kepala (per-head), pada 64, 128, atau 256 elemen. Tidak seperti norm RMS biasa yang beroperasi pada dimensi tersembunyi yang melebihi ribuan, yang memerlukan reduksi horizontal melalui memori bersama dan register, QK norm hanya memerlukan reduksi di seluruh warp, dengan kemungkinan memasukkan beberapa kepala ke dalam satu warp. Ini menjadi kasus penggunaan di mana kami dapat mengoptimalkan secara agresif pada waktu kompilasi untuk menemukan penetapan token-ke-warp terbaik untuk semua dimensi kepala dan tipe data.

Norm RMS adalah operasi yang cukup mudah, meskipun memerlukan penjumlahan horizontal untuk menghitung norm kuadrat dari sebuah vektor:

cpp

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

Karena kepala perhatian biasanya berukuran kecil dan operasi per-utas sangat bervektorisasi, satu warp dapat menghitung norm kuadrat. Untuk memaksimalkan bandwidth memori, setiap utas harus memproses 128 bit masukan, yang setara dengan 8 nilai bfloat16. Akibatnya, setengah-warp (half-warp) sudah cukup untuk memproses sebuah kepala, dan norm kuadrat dapat ditentukan menggunakan reduksi horizontal yang dilakukan oleh shuffle_sync_bfly. Berdasarkan dimensi kepala dan tipe data, kita dapat menentukan jumlah kepala per warp serta menentukan argumen dan jumlah putaran reduksi pada saat kompilasi.

Kita kemudian dapat membangun kernel khusus dengan baris kode yang relatif sedikit, dengan mengandalkan pemuatan/penyimpanan yang bervektorisasi secara otomatis dan menggunakan operasi spesifik arsitektur tingkat rendah untuk reduksi warp. Kernel menentukan kepala mana (Q atau K) yang akan dikerjakan dan mengidentifikasi bobot yang sesuai untuk diterapkan. Kernel memuat seluruh vektor dari kepala dan pertama-tama menjumlahkannya secara lokal. Selanjutnya, berdasarkan jumlah utas per kepala, norm kuadrat dihitung. Terakhir, norm RMS, yang bervektorisasi secara otomatis, ditentukan dan ditulis kembali ke memori. Parameter khusus model, seperti WEIGHT_BIAS dan EPS, dilipat secara konstan (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)

Dibandingkan dengan kernel norm RMS biasa, varian khusus ini memberikan throughput dan latensi 2-3x lebih baik untuk prefill maupun decode.

Pengiriman dan penggabungan MoE mengacak (shuffle) token masukan ke pakar (experts) tempat token tersebut diarahkan, melakukan akumulasi berbobot setelah kernel MoE GEMM melakukan perkalian matriks yang relevan. Performa menjadi hal yang sangat menarik, karena berbagai produk mengandalkan model MoE dengan berbagai ukuran. Berdasarkan jenis penerapannya, kami membedakan berbagai kernel untuk dioptimalkan:

TP=1: Kernel lokal untuk satu GPU tunggal. Pengiriman mengacak token ke pakar yang relevan dan penggabungan mengumpulkannya. Digunakan terutama untuk penerapan Qwen3 30B pada H100.

DP=1, TP ≠ 1: Kernel paralel-tensor, satu node (singlenode). Masukan direplikasi, pengiriman memilih token untuk peringkat saat ini, dan penggabungan mengumpulkannya di semua peringkat. Akumulasi melalui NVLink atau NVLS. Digunakan untuk prefiller berukuran besar.

DP ≠ 1, NVLink: Kernel paralel-data dengan paralelisme tensor opsional. Setiap peringkat mengirimkan token ke setiap peringkat lainnya yang menghosting pakar tertentu. Informasi perutean digabungkan, token ditransfer melalui NVIDIA NVLink. Digunakan untuk penerapan H200 atau GB200 paralel-data.

DP ≠ 1, InfiniBand: Kernel paralel-data Peer-to-Peer dan paralel-tensor yang beroperasi di seluruh NVLink dan NVIDIA Quantum InfiniBand. Digunakan untuk penerapan Hopper multinode.

Kernel lokal menghitung jumlah token yang ditetapkan ke setiap pakar berdasarkan tabel perutean dan menetapkan indeks unik ke setiap token dalam daftar token milik masing-masing pakar. Setelah langkah perutean, token disalin ke dalam buffer yang dapat diberi padding secara opsional untuk diteruskan ke kernel MoE GEMM. Dengan menggunakan lokasi token, kernel penggabungan kemudian melakukan akumulasi dalam register, dan menuliskan token akhir ke dalam buffer keluaran.

Fase pengiriman bervariasi antara prefill dan decode: untuk jumlah token yang sedikit, hanya ada sedikit informasi perutean yang dibaca, sehingga kami memprosesnya dalam satu blok dan memberi sinyal ke sisa kisi, dengan masing-masing blok menyalin potongan token. Penetapan token ditentukan dengan menaikkan penghitung secara atomik (atomically incrementing) di memori bersama. Karena prefill harus memproses lebih banyak informasi perutean, perutean dibagi ke dalam beberapa blok di kernel terpisah. Indeks pertama-tama ditetapkan secara lokal menggunakan penghitung di memori bersama dan kemudian digabungkan di semua blok di memori global untuk menentukan posisi setiap token. Kisi kernel ini dijepit (clamped) ke jumlah SM agar dapat menggunakan penghalang berbasis kisi (grid-wise barriers). Kernel berikutnya melakukan penyalinan, diluncurkan di sejumlah besar blok dan warp untuk memanfaatkan bandwidth memori yang tersedia secara penuh.

Kernel penggabungan relatif sederhana, karena melakukan akumulasi bervektorisasi, dengan mengandalkan informasi perutean dan offset yang dihitung oleh pengiriman. CuTeDSL terbukti berharga di sini, karena komputasi dapat diekspresikan melalui tensor CuTe sederhana yang dikompilasi menjadi kode yang efisien untuk semua tipe data dan dimensi tersembunyi.

Dalam penerapan paralel-tensor, informasi perutean dan token masukan direplikasi. Kernel pengiriman lokal dapat digunakan kembali sepenuhnya dengan sedikit perubahan untuk merutekan dan memilih hanya token yang ditetapkan ke pakar yang terletak di peringkat saat ini. Dengan beberapa parameter dan spesialisasi waktu kompilasi, kernel pengiriman lokal digunakan kembali.

Karena jumlah peringkat biasanya lebih rendah daripada jumlah pakar tempat token diarahkan, dalam kasus rata-rata dapat diasumsikan bahwa setiap token memiliki setidaknya satu replika yang diarahkan ke setiap peringkat. Dalam situasi ini, kernel penggabungan lokal dapat digunakan kembali untuk mengumpulkan salinan lokal dari suatu token, diikuti dengan all-reduce untuk menjumlahkan di seluruh peringkat yang berpartisipasi. Dalam kasus kernel prefill, kami meluncurkan kernel terpisah, namun decode menggabungkan kedua operasi tersebut menjadi satu peluncuran tunggal.

Untuk all-reduce, kami menggunakan NVLink biasa jika hanya dua rank yang terlibat, lalu beralih ke NVLS ketika 4 atau lebih rank terlibat untuk peningkatan throughput sebesar 10-20%. Berkat DSL, kami dapat membagikan sebagian besar logika reduksi dan pertukaran buffer peer antara MoE dan kernel all-reduce biasa. Ini adalah situasi di mana kontrol yang erat atas perangkat keras sangat bermanfaat, karena operasi DSL pada multimem.ld_reduce dan multimem.st diperlukan untuk NVLS. Implementasi reduksi itu sendiri sebagian besar mengikuti contoh all_reduce dari pengujian NCCL.

Kernel Paralel Data mengandalkan NVLink untuk menggabungkan token dari beberapa perangkat. Pada fase pengiriman, sebelum bertukar token, peer terlebih dahulu merutekan secara lokal untuk mengidentifikasi offset target pada setiap peringkat untuk setiap token yang mereka kirim. Selanjutnya, mereka mulai menyalin token ke dalam buffer penerima yang dialokasikan untuk setiap pengirim guna menghindari konflik. Penulisan rilis-simpan (release-store) per-token digunakan untuk memberi sinyal penyelesaian. Di sisi penerima, informasi perutean digabungkan agar dapat memadatkan semua token ke dalam satu buffer tunggal untuk MoE GEMM. Penerima memindai (poll) bendera yang ditulis peer dengan muatan-perolehan (acquire-load) untuk menunggu penulisan selesai.

Di sisi penggabungan (combine), token disalin melalui NVLink ke peringkat sumber, dengan bendera per-token yang menunjukkan tanda terima. Kernel penggabungan mirip dengan varian lokal dan singlenode, dengan sedikit perbedaan pada pengalamatan (addressing) dan beberapa logika sinkronisasi tambahan.

Kami membangun pustaka pembungkus PTX untuk mengekspos muatan dan penyimpanan rilis-perolehan (release-acquire loads and stores) yang dicakup pada tingkat sistem maupun perangkat guna mengimplementasikan kernel ini. Dengan meningkatkan granularitas bendera sinkronisasi dan mengalokasikan lebih banyak blok untuk mengirim dan menerima token melalui NVLink, kami mencapai peningkatan 2x lipat dibandingkan kernel CUDA kami sebelumnya.

Kernel Peer-to-Peer kami sebelumnya mencapai performa mutakhir pada GPU Hopper yang dipasangkan dengan adaptor InfiniBand ConnectX-7. Sejak saat itu, kami telah memindahkan kernel penggabungan ke CuTeDSL untuk mendapatkan peningkatan performa lebih lanjut. Karena versi CUDA sebelumnya sangat mengandalkan spesialisasi waktu kompilasi, versi tersebut terbukti menjadi hambatan waktu kompilasi di repositori kami, yang memaksa kami menghilangkan beberapa parameter. Dengan bermigrasi ke DSL yang dikompilasi secara JIT, kami melakukan spesialisasi untuk lebih banyak parameter dan menyesuaikan alokasi kisi, serta memperkenalkan perbedaan lebih lanjut antara prefill dan decode untuk peningkatan latensi sekitar 3%.

Dari perspektif kernel ROSE, dukungan Blackwell terbagi menjadi dua kategori: kernel yang agnostik terhadap arsitektur (architecture-agnostic) dan yang spesifik terhadap arsitektur (architecture-specific). Kernel yang spesifik terhadap arsitektur, yang mengandalkan inti Tensor, memerlukan pekerjaan eksplisit untuk beradaptasi. Namun, berkat penggunaan DSL, kernel yang agnostik terhadap arsitektur dapat beradaptasi dengan mudah.

GPU Blackwell memiliki hingga 160 SM, kontras dengan Hopper yang memiliki 132 SM. Bagi sebagian besar kernel, ini adalah konstanta waktu kompilasi dan logika dalam kode peluncuran di sisi host memperhitungkannya untuk menentukan kernel optimal (prefill vs decode) yang akan dijalankan. Selain itu, Blackwell juga memperkenalkan instruksi muat 256-bit, yang dapat dipancarkan dengan mudah melalui peningkatan lebar vektor (dimensi terakhir dari tensor CuTe) untuk meningkatkan pemanfaatan bandwidth memori.

Blackwell memperkenalkan dukungan perangkat keras untuk sejumlah besar skema kuantisasi (MXFP8, MXFP4, NVFP4). Dukungan untuk skema ini diimplementasikan terutama dengan menyesuaikan kernel yang ada untuk penskalaan berbasis blok FP8. Dukungan JIT sangat penting di sini, karena mode baru dapat didukung tanpa membuat ruang konfigurasi waktu kompilasi membengkak. Di masa lalu, kami mendapati Triton canggung dalam mengekspresikan reduksi horizontal di beberapa elemen saja, yaitu di atas setengah atau seperempat warp. Kontrol yang lebih detail atas reduksi sangat penting untuk mencapai performa optimal saat menghitung nilai maksimum absolut di 16 atau 32 elemen masing-masing untuk NVFP4 dan MXFP8/MXFP4.

Sebagai bahasa yang relatif baru, LLM memiliki sedikit paparan terhadap kode CuTeDSL selama pelatihan awal (pre-training) mereka. Selain itu, gaya kode kami saat ini, dengan kernel yang dioptimalkan untuk aktivasi dan norm, sedikit berbeda dari contoh CUTLASS atau Quack yang sudah ada, yang sangat dioptimalkan untuk TMA dan instruksi perkalian matriks spesifik arsitektur. Mencapai titik di mana LLM dapat membantu dalam penulisan kernel merupakan proses yang bertahap.

Kami secara manual membangun kumpulan awal korpus kernel dan fungsi utilitas, dengan mengikuti pola yang konsisten, sebelum kami dapat mengandalkan agen untuk konversi kernel. Seiring berkembangnya korpus tersebut, kami dapat melakukan terjemahan kasar dari CUDA ke CuTeDSL, dengan pengerjaan ulang manual yang signifikan setelahnya. Khususnya, kami mendapati bahwa model mengalami kesulitan dengan konversi pointer-ke-tensor dan vektorisasi, serta terjebak saat mencoba memperbaiki kesalahan tipe ketika bertujuan agar sebuah tes dapat dikompilasi. Model-model tersebut menunjukkan sedikit pemahaman tentang sifat berbasis kisi dari kernel, sehingga mencampuradukkan sinkronisasi utas dan kisi. Yang paling membuat frustrasi, ada godaan untuk berhalusinasi membuat terjemahan palsu dari pustaka khusus CUDA alih-alih membangun berdasarkan pengganti CuTeDSL kami sendiri. Seiring waktu, seiring berkembangnya korpus kernel CuTeDSL, kami mendapati bahwa LLM dapat membangun sebagian besar kode kerangka (skeleton code) dan mengisi bagian yang kosong, sehingga memungkinkan kami untuk fokus pada desain tingkat tinggi.

Dibandingkan dengan CUDA, kami mendapati bahwa CuTeDSL memberikan performa waktu kompilasi yang jauh lebih baik, memungkinkan kami melakukan spesialisasi untuk lebih banyak parameter dan mencapai peningkatan performa. Meskipun kernel sedikit lebih bertele-tele karena kurangnya abstraksi tertentu dibandingkan dengan Triton, kami mendapati bahwa kompleksitas tersebut dibenarkan karena kontrol yang lebih detail atas perangkat keras memungkinkan kami membangun kernel yang lebih cepat.

Pada saat penulisan ini, kesenjangan performa antara kernel CUTLASS C++ dan CuTeDSL sangat minimal, membuka peluang menarik untuk menyusun dan menggabungkan kernel GEMM berperforma tinggi dengan komunikasi antarperangkat, normalisasi, atau aktivasi. Dengan berkomitmen pada DSL ini, kami akan dapat mengekstrak performa yang lebih tinggi dari perangkat kami guna meningkatkan throughput di seluruh beban kerja kami, mulai dari penyematan hingga LLM.

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

Inferensi di Perplexity AI

Mengapa CuTeDSL?

Pengalaman pengembang

Ruang konfigurasi yang besar

Spesialisasi prapengisian dan dekode

Sinkronisasi kisi (Grid synchronization)

Kernel CuTeDSL

QK Norm

Pengiriman/Penggabungan Mixture-of-Experts (MoE Dispatch/Combine)

Lokal (Local)

Satu Node (Singlenode)

Paralel Data (Data Parallel)

Peer-to-Peer

Dukungan Blackwell (Blackwell Support)

Generasi Kernel Berbantuan AI (AI-Assisted Kernel Generation)

Kesimpulan dan Pekerjaan Mendatang

Referensi

Kutipan