Cara membuat GPU AMD berjalan dengan kecepatan ‘brrr’

• HipKittens adalah kumpulan primitif pemrograman yang dirancang untuk mengeluarkan potensi performa GPU AMD, dengan mengoptimalkan akses memori, penjadwalan, dan penggunaan ulang cache
• GPU AMD MI355X memiliki struktur 256 compute unit dan 8 chiplet (XCD), serta menyediakan register file besar dan instruksi matrix core yang halus
• Tidak seperti NVIDIA, AMD tidak memiliki reallocasi register, instruksi matriks asinkron, dan mbarrier, sehingga penjadwalan 8-wave ping-pong dan 4-wave interleave lebih efektif dibanding wave specialization
• HipKittens meningkatkan lokalitas cache L2 dan LLC melalui penjadwalan grid yang sadar chiplet, dan mencapai peningkatan bandwidth maksimum serta TFLOPS pada operasi GEMM dan Attention
• Pendekatan ini melengkapi kurangnya kematangan perangkat lunak di ekosistem GPU AMD dan menyediakan fondasi untuk meningkatkan skalabilitas komputasi AI berbasis beragam perangkat keras

---

Arsitektur dan karakteristik performa GPU AMD CDNA

• GPU AMD MI355X mencakup 256 compute unit (CU), dan tiap CU terdiri dari 4 SIMD
  • Satu SIMD mengeksekusi wave yang terdiri dari 64 thread, berbeda dengan warp 32-thread milik NVIDIA
• MI355X memiliki SRAM setara 70% dari B200 (165KB), dan tidak memiliki fitur instruksi perkalian matriks asinkron, reallocasi register, akselerasi tensor memory, serta mbarrier
• Sebaliknya, ia menyediakan register file 2 kali lebih besar dan jumlah prosesor 60% lebih banyak (256 CU vs 160 SM)
  • Mendukung instruksi matrix core kecil dan terperinci, serta memiliki kemampuan load langsung global→shared memory (mirip TMA)
• AMD mengadopsi arsitektur chiplet yang terdiri dari 8 chiplet (XCD); setiap XCD memiliki cache L2 independen, dengan cache LLC di tingkat atas
• Menurut tabel, MI355X memiliki performa komputasi BF16 2.5 PFLOPs, MXFP8 5.0 PFLOPs, MXFP6 10.1 PFLOPs, serta kapasitas memori 288GB dan bandwidth 8TB/s

Tantangan desain kernel untuk AMD

• Optimasi akses memori: karena kendala kompiler HIPCC dan perilaku I/O yang tidak dipublikasikan, desain pola tata letak data dan swizzle menjadi penting
• Penjadwalan intra-prosesor: AMD perlu memanfaatkan register file dan instruksi matriks kecil, bukan shared memory
• Penjadwalan antar-prosesor: karena struktur berbasis chiplet, distribusi kerja perlu mempertimbangkan efek NUMA pada level cache

Pola akses memori HipKittens

• HipKittens(HK) menggunakan tile sebagai unit data dasar, dan menyediakan fungsi operasi mirip PyTorch
  • Tile didefinisikan oleh tipe data, ukuran, dan layout, serta menangani beragam input melalui metaprogramming template C++
• Penjadwalan register: HIPCC tidak dapat menggunakan register tertentu sebagai input MFMA, sehingga HK menyediakan fitur pinning register eksplisit
  • Pengembang bisa menetapkan register secara langsung untuk menulis kernel dengan performa maksimum
• Layout register: di AMD, layout berbeda menurut tipe data dan bentuk matriks, sehingga satu pola swizzle tunggal tidak memungkinkan
  • Sebagai contoh, tile bf16 16×16 dan tile bf16 16×32 memerlukan pola swizzle yang berbeda
• Struktur fase instruksi: instruksi shared memory AMD memiliki grup fase yang tidak berurutan, dengan dokumentasi internal yang minim
  • HK menyediakan solver hasil reverse engineering untuk ini
• Pembuatan alamat: AMD mendukung load HBM→shared memory secara asinkron, dan melakukan optimasi dengan swizzle alamat HBM

Penjadwalan intra-prosesor: pola wave

• Wave specialization efektif di NVIDIA, tetapi pada AMD performanya menurun karena tidak adanya reallocasi register
  • Wave producer menempati register yang tidak diperlukan, sementara wave consumer mengalami spill karena kekurangan register
• Hasil eksperimen HK menunjukkan bahwa wave specialization di AMD menyebabkan penurunan intensitas aritmetika dan bottleneck memori
  • Contoh: pada GEMM, konfigurasi HK 0/8 mencapai 1605 TFLOPs, sedangkan CUTLASS 1570 TFLOPs
• Pola penjadwalan alternatif
  • 8-wave ping-pong: dua wave secara bergantian menjalankan klaster memori/komputasi
  • 4-wave interleave: satu wave menjalankan interleaving halus antara memori dan komputasi
  • 8-wave memiliki kode yang ringkas, sedangkan 4-wave lebih terperinci tetapi kodenya lebih panjang
  • Pada GEMM dan Attention Forward, 8-wave mencapai performa setara SoTA

Penjadwalan antar-prosesor: pendekatan sadar chiplet

• AMD MI355X memiliki 8 chiplet XCD, dan tiap chiplet memiliki cache L2 independen
  • Thread block dialokasikan ke chiplet dengan cara round-robin, sehingga urutan grid langsung memengaruhi efisiensi penggunaan ulang cache
• Tata letak row-major yang sederhana menghasilkan tingkat penggunaan ulang cache L2 yang rendah dan menyebabkan kehilangan bandwidth
  • Contoh: L2 55%, LLC 95%, 15.1 TB/s, 1113 TFLOPs
• HK memperkenalkan penjadwalan grid yang sadar chiplet, memanfaatkan sekaligus lokalitas cache L2 dan LLC
  • Thread block dikelompokkan berdasarkan area berdekatan pada matriks output untuk memaksimalkan penggunaan ulang data input

Contoh kernel nyata

• Hot loop pada kernel Attention Forward dan BF16 GEMM menggunakan skema 8-wave ping-pong dari HK
  • Tiap loop secara bergantian mengeksekusi klaster Compute–Memory dan disinkronkan dengan schedule barrier
  • Dalam contoh kode, operasi HK seperti mma_AtB, load, exp2, col_sum digunakan berulang kali

Kesimpulan: AMD di era AI multi-silicon

• HipKittens mencapai performa yang kompetitif pada AMD CDNA3 dan CDNA4
  • Tiga inti utamanya: akses memori teroptimasi, penjadwalan wave berfokus AMD, dan penjadwalan grid sadar chiplet
• Kernel HK mencapai performa tertinggi menurut standar AMD dan juga kompetitif dengan kernel NVIDIA Blackwell
• Demi keberagaman komputasi AI, perlu ada perluasan aksesibilitas GPU AMD, dan HipKittens menyediakan fondasi perangkat lunak inti untuk itu
• Peningkatan penjadwalan register HIPCC dari AMD disebut sebagai area pengembangan penting ke depan