Jantung AI Berdetak Lebih Cepat: Membedah Otot nn.Linear dan Fused MLP, Mending Tangan Manusia atau Otak Kompilator?

Di dunia AI yang serba “ngebut” ini, seringkali kita terjebak pada ilusi bahwa algoritma canggih secara otomatis berarti performa super kilat. Padahal, di balik setiap model AI yang “pintar”, ada kerja keras para “majikan” (manusia) yang telaten mengoptimalkan setiap jengkal kode, bahkan hingga ke urat nadi GPU.

Dalam seri kedua “Profiling di PyTorch” ini, kita akan membongkar lebih dalam bagaimana memerintah AI agar bekerja lebih efisien. Kita mulai dari blok bangunan fundamental: nn.Linear, lalu meningkat ke Multilayer Perceptron (MLP) yang menjadi tulang punggung banyak model modern. Apakah kompilator AI sudah cukup cerdas, atau sentuhan tangan manusia masih tak tergantikan?

Sebelum kita mulai, ingat dua prinsip emas ini:

1. Sebuah kernel GPU adalah program mini yang berjalan paralel di ribuan thread GPU.
2. CPU menjadwalkan dan meluncurkan kernel-kernel ini. Sebagian besar overhead PyTorch yang terlihat di trace profiler adalah pekerjaan penjadwalan ini.

Dari Matmul-Add ke nn.Linear: Evolusi yang ‘Bikin Malas’

nn.Linear itu sejatinya hanya bungkus manis dari operasi perkalian matriks (matmul) dan penambahan (add) yang sudah kita bahas sebelumnya. Bedanya, dia punya “otot” (weight dan bias) sendiri dan metode forward yang ramah pengguna. Namun, jangan salah sangka, ada trik kecil yang bikin dia lebih “malas” (baca: efisien).

Kalau kita intip jejaknya pakai profiler, kita akan lihat ada operasi aten::t (transpose) sebelum aten::addmm. Ini bukan sihir! nn.Linear hanya membalik orientasi matriks bobotnya di CPU, tapi tanpa memindahkan data sungguhan di GPU. Ibaratnya, dia cuma ganti “label” arah tanpa perlu bongkar muat koper. Ini efisiensi ala robot yang cuma ganti pikiran, bukan gerak fisik.

Kenapa Tidak Ada Kernel Penambahan Terpisah? Ini Dia Epilog Rahasia AI!

Hebatnya, kita tidak akan menemukan aten::add (penambahan bias) terpisah di rantai kerja nn.Linear. Kenapa? Karena penambahan bias itu sudah “dilipat” ke dalam kernel perkalian matriks, menggunakan teknik yang disebut epilog.

Epilog itu seperti “pencuci mulut” di akhir makan berat. Komputasi kecil yang dilakukan kernel GEMM (General Matrix Multiply) tepat sebelum hasilnya ditulis ke HBM (memori utama GPU). Menambahkan bias, mengaplikasikan aktivasi, atau menskalakan konstanta, semuanya adalah epilog klasik. Tujuannya sederhana: hindari bolak-balik memori yang bikin boros waktu dan energi. Ini bukti bahwa AI, bahkan di level dasar, sudah lumayan “cerdas” dalam manajemen tugas.

Menariknya, kernel addmm yang kita lihat di Bagian 1 saat menggunakan --compile, sebenarnya sudah dipakai oleh nn.Linear secara eager. Jadi, untuk satu nn.Linear, torch.compile tidak punya banyak “sihir” lagi untuk digabungkan. Ini seperti menyuruh asisten rumah tangga yang sudah rajin untuk “lebih rajin lagi” padahal dia sudah maksimal. Robot itu patuh, tapi kadang kurang inisiatif kalau tugasnya cuma satu.

Baca juga artikel menarik lainnya di kategori Update Algoritma.

Akal Manusia Vs. Otak Kompilator: Siapa yang Lebih Cerdas?

Saat kita “mengkompilasi” satu lapisan nn.Linear dengan torch.compile, kita tidak melihat perubahan signifikan pada kernel GPU. Kernel yang berjalan tetap sama, yaitu cuBLAS GEMM yang sudah ada. Namun, yang hilang adalah overhead CPU untuk penjadwalan pandangan matriks (seperti aten::t). Kompilator menelusuri rantai pandangan ini, menghitung langkah-langkah yang diperlukan, dan langsung memanggil aten::addmm dengan parameter yang sudah “dimasak”. Beberapa mikrodetik kerja CPU hilang, sementara GPU melakukan pekerjaan identik.

Ini penting untuk dicerna: Jangan asal pakai torch.compile setiap kali model terasa lambat. Untuk satu operasi GEMM dengan bias, kompilator AI tidak punya banyak lagi yang bisa dioptimalkan. Ia butuh lebih dari satu operasi untuk bisa “menggabungkan” sesuatu, layaknya asisten yang baru akan berinisiatif menggabungkan tugas jika ia melihat ada beberapa tugas yang mirip dan bisa dilakukan sekaligus. Jadi, AI ini pintar, tapi *masih perlu sekolah* untuk melihat gambaran besar jika hanya diberi tugas sepotong-sepotong.

MLP: Ketika Linear Berbaris dan Kompilator Mulai Beraksi

Sekarang, mari kita tumpuk tiga lapisan nn.Linear menjadi Multilayer Perceptron (MLP) dengan aktivasi GeGLU. Di mode eager, kita akan melihat lima kernel GPU berjalan: tiga untuk nn.Linear, satu untuk GeLU, dan satu untuk perkalian. Setiap kernel ini menulis hasilnya ke HBM, lalu kernel berikutnya membacanya lagi. Ini seperti mengambil barang dari lemari, meletakkannya di meja, lalu mengambilnya lagi untuk ditaruh di tempat lain. Boros waktu!

Di sinilah torch.compile mulai menunjukkan taringnya. Dua kernel pointwise (GeLU dan perkalian), ditambah operasi reshape, dilebur menjadi satu kernel Triton yang terpadu, bernama triton_poi_fused__unsafe_view_gelu_mul_0. Apa artinya? Data menengah tidak perlu lagi bolak-balik ke HBM; semuanya diproses langsung di register GPU. Ini ibaratnya asisten yang kini bisa mengambil barang dari lemari, memprosesnya di tempat, dan langsung menaruhnya di tempat tujuan tanpa perantara. Efisien!

Ini adalah bukti bahwa ketika ada lebih banyak operasi untuk digabungkan, kompilator AI bisa jadi sangat efektif. Ia mengoptimalkan alur data, mengurangi lalu lintas memori, dan membuat jantung AI berdetak lebih cepat. Namun, ia punya kelemahan: ia mengkhususkan diri untuk bentuk input statis. Jika ukuran batch, panjang urutan, atau dimensi tersembunyi berubah, kompilator harus “menghitung ulang” semuanya dari awal. Chip AI seperti Microsoft Maia 200 terus berinovasi di sisi hardware untuk mendukung komputasi seperti ini, namun batasan kompilator tetap ada.

Memilih Jalan Ninja: Kernel yang Disetel Tangan Manusia

Kita sudah melihat PyTorch eager dan torch.compile memilih kernel. Sekarang, bagaimana jika kita memakai kernel yang ditulis dan disetel sendiri oleh manusia ahli? Dengan bantuan library kernels dari Hugging Face Hub, kita bisa dengan mudah menggunakan lapisan LigerGEGLUMLP.

Keunggulan kernel buatan tangan ini ada dua:

1. Fusi yang sudah “dibakar”: Seperti kompilator, kernel buatan tangan ini juga menggabungkan operasi GeLU dan perkalian menjadi satu kernel. Bedanya, ini dilakukan tanpa perlu kompilator. Tidak ada guard Dynamo, tidak ada latensi kompilasi, dan tidak ada risiko kompilasi ulang. Ini seperti memiliki resep rahasia yang sudah teruji dan tidak perlu dimasak ulang setiap kali ada bahan yang sedikit berbeda.
2. Parameter peluncuran yang disetel untuk hardware: Kernel ini tidak menebak ukuran blok secara acak. Parameter peluncurannya sudah dipilih dengan cermat untuk arsitektur GPU spesifik. Ini membuahkan performa yang sangat stabil. Bahkan, di beberapa kasus, kernel buatan tangan ini bisa lebih cepat, atau setidaknya sekompetitif kernel hasil kompilasi.

Perlu diingat, ada trade-off di sini. Kernel Liger mungkin sedikit lebih lambat dari kernel hasil kompilasi yang sangat spesifik untuk satu bentuk input. Namun, kernel Liger memberikan fleksibilitas. Ia bisa berjalan di berbagai bentuk input tanpa perlu kompilasi ulang. Ini seperti tukang bangunan yang punya “rasa” dan bisa menyesuaikan diri dengan berbagai kondisi, tidak terpaku pada satu desain saja. Sementara itu, di ranah lain, bahkan AMD pun mulai membuktikan taringnya dalam mengoptimalkan AI tanpa CUDA, menunjukkan bahwa pertarungan performa ini jauh dari kata usai.

Jadi, pilih mana? Kecepatan mutlak dengan risiko kompilasi ulang setiap ada perubahan (ala kompilator), atau fleksibilitas dan stabilitas dengan sedikit kompromi kecepatan (ala tangan manusia)? Itu adalah keputusan “majikan” yang harus dibuat berdasarkan kebutuhan nyata.

Kesimpulan: Robot Memang Cerdas, Tapi Akal Manusia Tetap di Atas Segalanya

Dari penjelajahan kita, terlihat jelas bahwa optimasi performa AI adalah seni yang menggabungkan kecerdasan sistem dan sentuhan manusia. PyTorch sudah melakukan banyak hal di balik layar dengan fusi epilog. Kompilator seperti torch.compile mampu melebur operasi pointwise menjadi kernel tunggal, mengurangi beban memori. Namun, keahlian manusia dalam menulis dan menyetel kernel masih relevan, terutama untuk fleksibilitas dan menghindari overhead kompilasi.

Pada akhirnya, seberapa pun canggihnya AI, ia tetaplah alat yang membutuhkan “majikan” yang cerdas dan teliti. Tanpa akal sehat, kemampuan menganalisis, dan keputusan strategis dari manusia, AI hanyalah tumpukan kode yang menunggu perintah. Kita harus bisa menebak apa yang akan dilakukan robot, membandingkannya dengan realita, dan memperbaiki jika ada yang melenceng. Ini adalah esensi menjadi “Majikan AI”.

Ingat, secepat-cepatnya AI memproses data, lebih cepat lagi kamu harus menjelaskan kenapa tagihan listrik GPU bulan ini membengkak.

Sumber Berita

Artikel ini dirangkum dari sumber asli di “Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP”

Gambar oleh: Hugging Face