Optimisasi FlashAttention-3/4 tidak meningkatkan kinerja GPU RTX

Eksplorasi penerapan optimisasi FlashAttention-3 dan FlashAttention-4 pada GPU RTX konsumen mengungkapkan bahwa peningkatan yang berfokus pada datacenter ini tidak memberikan gain kinerja. Penulis membangun ulang kernel perhatian dari awal untuk menguji apakah teknik FA-3/4 dapat mengungguli fallback FA-2 yang ada pada perangkat keras konsumen.

• Instruksi tensor-core yang lebih cepat (WGMMA) yang menjadi inti FA-3 tidak tersedia pada GPU RTX.
• Tensor Memory Accelerator (TMA) hadir di sm_120 tetapi tidak membantu karena transportasi bukan bottleneck.
• Spesialisasi warp menghasilkan kinerja negatif, meningkatkan latensi dari 206us menjadi 213us.
• Optimisasi eksponensiasi simulasi FA-4 tidak relevan karena RTX 5090 tetap terikat tensor-core daripada SFU.

Studi ini menyimpulkan bahwa FlashAttention-2 mewakili batas kinerja untuk perhatian reguler pada silikon konsumen, karena peningkatan lebih lanjut akan memerlukan pengorbanan akurasi untuk presisi yang lebih rendah.

Analisis ini berfokus pada rezim prefill dan compute-bound; decoding terhadap cache KV besar tetap menjadi tantangan memory-bound terpisah di mana teknik split-KV lebih relevan.