Daily Akamuga – Para peneliti dari University of California, Berkeley, dan University of Texas, Austin, telah meluncurkan sebuah perpustakaan baru bernama Flash-KMeans. Library ini merupakan implementasi dari algoritma k-means yang telah ada sejak lama, namun dengan penekanan pada pengolahan data yang lebih cepat saat digunakan dalam alur kerja kecerdasan buatan modern. Dengan meningkatnya frekuensi pemanggilan algoritma ini dalam siklus latihan dan inferensi, latensi per panggilan menjadi sangat penting.
Flash-KMeans adalah implementasi yang sadar input-output (IO) dari algoritma k-means standar yang dikenal sebagai Lloyd’s k-means. Perpustakaan ini tidak mengubah matematika dasar dari algoritma tersebut, melainkan memperbaiki cara data dipindahkan di GPU. Menurut laporan tim peneliti, menggunakan NVIDIA H200, Flash-KMeans dapat meningkatkan kecepatan hingga 17.9 kali lipat dibandingkan dengan baseline terbaik yang ada. Dibandingkan dengan cuML dari NVIDIA, kecepatan meningkat hingga 33 kali, dan hingga lebih dari 200 kali jika dibandingkan dengan FAISS, sebuah perpustakaan yang juga banyak digunakan di industri.
Inovasi dalam Flash-KMeans
Flash-KMeans ditulis dalam kernel Triton GPU dan tersedia di bawah lisensi Apache 2.0, dengan dapat diinstal menggunakan perintah pip. Hasil output dari algoritma ini identik secara matematis dengan k-means standar Lloyd. Kecepatan yang diperoleh berasal dari pengolahan aliran data pada tingkat kernel, bukan dari pengurangan pekerjaan yang dilakukan.
Dalam setiap iterasi standar Lloyd, terdapat dua tahap utama: tahap penugasan dan pembaruan centroid. Pada tahap penugasan, setiap titik dihitung jaraknya ke setiap centroid, dan kemudian titik terdekat dipilih. Sementara pada tahap pembaruan, rata-rata titik dalam setiap cluster dihitung untuk membentuk centroid baru. Keduanya biasanya terbatas pada memori, bukan kemampuan komputasi, sehingga menjadi hambatan utama ketika dijalankan di GPU.
Dua Hambatan yang Diterapkan
Dua hambatan utama yang menjadi fokus dalam pengembangan Flash-KMeans adalah:
-
Tahap Penugasan: Pada implementasi standar, sebuah matriks jarak D dibangun dalam memori bandwidth tinggi (HBM), yang kemudian dibaca kembali untuk mendapatkan nilai minimum. Proses ini pada N=65536 dan K=1024 memerlukan waktu yang signifikan. Flash-KMeans menggunakan metode yang disebut FlashAssign, yang membagi data menjadi potongan-potongan kecil untuk meminimalkan kompleksitas IO dari O(NK) menjadi O(Nd + Kd).
-
Tahap Pembaruan Centroid: Dalam teknik standar, penggunaan penambahan atomik dapat menyebabkan hambatan jika beberapa thread yang sama mencoba mengakses cluster yang sama secara bersamaan. Flash-KMeans memperkenalkan metode baru bernama Sort-Inverse Update, di mana ID cluster diurutkan terlebih dahulu untuk mengurangi konflik dan meningkatkan efisiensi pemrosesan.
Pencapaian Benchmark dan Evaluasi Kinerja
Tim peneliti melakukan pengujian pada NVIDIA H200 menggunakan CUDA 12.8. Mereka melakukan variasi pada ukuran N, K, dan batch size B, serta membandingkan performa dengan empat baseline yang telah dioptimalkan. Hasil pengujian menunjukkan bahwa Flash-KMeans mampu menyelesaikan tugas di luar memori, dengan satu miliar titik data dalam 41.4 detik, jauh lebih cepat dibandingkan dengan baseline lainnya.
Salah satu keunggulan dari Flash-KMeans adalah kemampuannya untuk menangani data dalam skala besar tanpa mengalami masalah memori yang sering terjadi pada implementasi standar. Dengan menggunakan teknik pemisahan aliran data, Flash-KMeans dapat menyelesaikan tugas secara lebih efisien di lingkungan produksi.
Potensi Penggunaan dan Aplikasi Masa Depan
Dengan kecepatan yang lebih tinggi untuk k-means yang tepat, Flash-KMeans membuka berbagai kemungkinan penggunaan, baik dalam konteks offline maupun online. Beberapa aplikasi potensial termasuk pengindeksan pencarian vektor, rute perhatian yang jarang, dan kompresi cache KV.
Dari segi antarmuka, penggunaan Flash-KMeans cukup sederhana. API-nya sangat mirip dengan faiss dan sklearn, memungkinkan integrasi yang mudah ke dalam alur kerja yang sudah ada. Perpustakaan ini memastikan bahwa pengguna dapat dengan mudah mengadopsi teknologi baru ini tanpa memerlukan pelatihan yang intensif.
Kesimpulan
Secara keseluruhan, Flash-KMeans memperbaharui cara algoritma k-means digunakan dalam praktik dengan memanfaatkan arsitektur GPU modern. Dengan mengatasi dua bottleneck utama dalam algoritma, Flash-KMeans tidak hanya meningkatkan kecepatan tetapi juga skala yang dapat ditangani dalam produksi. Dengan klaim kecepatan hingga 17.9 kali lipat dibandingkan dengan baseline terbaik, proyek ini menawarkan harapan besar untuk pengolahan data skala besar dan aplikasi kecerdasan buatan di masa mendatang.
