Profil PyTorch dari satu nn.Linear sampai MLP yang ter-fuse

Aritra Roy Gosthipaty bersama tim Hugging Face merilis lanjutan tutorial profiler PyTorch mereka, berjalan dari satu nn.Linear sampai MLP utuh dan menunjukkan letak keuntungan performa yang sebenarnya mudah dipetik. Temuan pertama: nn.Linear yang punya bias tidak butuh kernel add terpisah. PyTorch melipat bias ke epilogue GEMM, jadi penjumlahannya terjadi di kernel yang sama persis sebelum hasilnya dituliskan ke memori. Dispatcher GEMM lalu memilih dari binari kernel precompiled berdasar layout tensor, dengan nama seperti suffix tn yang mengkodekan pola stride. Komputasi logis yang sama bisa mendarat ke kernel berbeda tergantung input shape, dan satu-satunya cara melihatnya adalah lewat trace.

Contoh MLP-nya, blok feed-forward dengan aktivasi GeGLU, berjalan sebagai lima kernel GPU dalam mode eager: tiga matmul plus GeLU dan elementwise multiply yang terpisah. Memanggil torch.compile melipat GeLU dan multiply jadi satu kernel, menjaga tensor intermediate di register, dan menghemat sekitar 50 MB traffic memori per forward pass. Tulisan ini juga membandingkannya dengan library Liger kernel, yang menyediakan kernel hand-tuned dan version-pinned. Pada shape statis Liger ada di 92,8 µs vs 89,4 µs untuk versi compiled, sedikit lebih lambat, tapi tidak perlu recompile saat dimensi input berubah. Catatannya tertib soal metodologi dengan cara yang jarang muncul di tutorial profiler: bentuk dulu prediksi, baru baca trace, lalu gali hanya pada bagian yang tidak cocok antara prediksi dan trace.

Kenapa ini penting

Kalau kalian melatih atau menyajikan transformer dan belum pernah membuka trace profiler, ini pengenalan terbersih untuk ke titik mana waktu dan memori sebenarnya pergi di dalam satu blok MLP, dengan angka konkret yang menempel. Perbandingan torch.compile lawan Liger juga sekaligus jadi keputusan deployment: kalau shape serving kalian statis, compile lebih cepat; kalau shape sering berubah, kernel prebuilt menyelamatkan kalian dari biaya recompile.