Grouped-query attention (GQA) mengubah satu rasio struktural di dalam layer attention: jumlah query head dapat lebih besar daripada jumlah key dan value head. Beberapa query head memakai key dan value head hasil proyeksi yang sama. Komputasi pada sisi query tetap terpisah per head, sedangkan state KV dibagi di dalam setiap grup.
Asimetri ini berpengaruh langsung pada decoding autoregresif karena key dan value dari token sebelumnya disimpan di cache. Semakin sedikit KV head yang berbeda, semakin sedikit pula state KV per token yang perlu dipertahankan untuk langkah decoding berikutnya.
Jumlah head menentukan batas grup
Misalkan sebuah layer attention memiliki (H_q) query head dan (H_{kv}) key-value head, dengan (H_q > H_{kv}). Pada pengelompokan reguler ketika (H_q) habis dibagi (H_{kv}), setiap KV head melayani
[ G = \frac{H_q}{H_{kv}} ]
query head.
Untuk query head (h), implementasi dapat memasangkannya dengan KV head
[ g(h)=\left\lfloor \frac{h}{G}\right\rfloor. ]
Operasi attention untuk query head tersebut memakai query vector miliknya sendiri bersama key dan value vector dari grup (g(h)):
[ A_h=\operatorname{softmax}\left(\frac{Q_h K_{g(h)}^\top}{\sqrt{d}}\right), \qquad O_h=A_hV_{g(h)}. ]
Aturan berbagi ini merupakan bagian arsitektur, bukan sekadar optimisasi penyimpanan. Model menghasilkan lebih sedikit proyeksi key dan value yang berbeda dibandingkan multi-head attention yang memiliki satu KV head untuk setiap query head.
GQA berada di antara MHA dan MQA
Multi-head attention (MHA) umumnya memakai jumlah head yang sama untuk query, key, dan value. Pada ujung lain, multi-query attention (MQA) memakai banyak query head dengan satu key-value head bersama.
GQA menempati konfigurasi di antara keduanya. Jumlah KV head lebih dari satu, tetapi tetap lebih sedikit daripada query head. Paper GQA tahun 2023 dari Ainslie dan kolega mendeskripsikan susunan ini sebagai generalisasi MQA dengan jumlah KV head di antara MHA dan MQA.
Hubungan ketiganya dapat ditulis ringkas:
[ H_{kv}=H_q \quad \text{untuk MHA}, ]
[ 1 < H_{kv} < H_q \quad \text{untuk GQA}, ]
[ H_{kv}=1 \quad \text{untuk MQA}. ]
Persamaan tersebut hanya menyatakan pola berbagi head. Persamaan itu tidak menyatakan kualitas model, latensi, atau perilaku memori yang identik pada setiap implementasi.
Ukuran KV cache mengikuti jumlah KV head
Pada decoding autoregresif dengan cache, setiap layer menyimpan key dan value vector untuk posisi token sebelumnya. Jika metadata, alignment, detail kuantisasi, dan padding khusus implementasi diabaikan, jumlah scalar yang tersimpan untuk satu layer berskala sebagai
[ 2 \times T \times H_{kv} \times d, ]
dengan (T) sebagai panjang sequence yang tersimpan dan (d) sebagai dimensi tiap KV head. Faktor dua berasal dari penyimpanan key dan value.
Jika dua arsitektur memakai (T), (d), format elemen, dan jumlah layer yang sama, penurunan (H_{kv}) menurunkan jumlah scalar KV persisten secara proporsional. Ini merupakan efek memori langsung dari berbagi KV head.
Batas klaim tersebut tidak mencakup seluruh memori serving. Runtime juga menyimpan bobot model, buffer attention sementara, overhead allocator, metadata request, dan state lain. Kernel tertentu dapat membuat view yang diperluas atau repetisi sementara saat komputasi walaupun cache persisten tetap memakai representasi KV yang ringkas.
Berbagi KV tidak menyamakan query head
Query head dalam grup yang sama tetap dapat memakai proyeksi query berbeda. Karena itu, masing-masing dapat menghasilkan logit attention dan distribusi probabilitas posisi token yang berbeda saat membaca key bersama.
Untuk dua query head (a) dan (b) yang dipetakan ke grup KV (j),
[ Q_a \neq Q_b ]
dapat menghasilkan
[ Q_aK_j^\top \neq Q_bK_j^\top. ]
Proyeksi KV bersama membatasi representasi key-value yang tersedia bagi query head tersebut, tetapi tidak menyatukan query vector menjadi satu head. Output keduanya masih dapat berbeda meskipun sama-sama membaca (K_j) dan (V_j).
Perbedaan ini penting saat memeriksa bentuk tensor. Memperluas tensor KV ringkas agar sesuai dengan grup query pada antarmuka kernel tidak menciptakan KV head baru dengan proyeksi independen. Data tersebut dapat berupa view atau repetisi dari grup yang sama.
Parameter proyeksi ikut berubah
Berbagi KV head tidak hanya memengaruhi cache decoding. Layout proyeksi yang menghasilkan lebih sedikit key dan value head memerlukan lebih sedikit channel output untuk proyeksi K dan V jika dimensi per head tetap sama.
Jika lebar model yang masuk ke proyeksi adalah (D), proyeksi dense K dan V yang terpisah memiliki lebar output (H_{kv}d). Jumlah parameternya ikut berskala terhadap (H_{kv}). Proyeksi Q dapat tetap memiliki lebar (H_qd).
Layout parameter aktual berbeda antarimplementasi. Sebagian runtime menggabungkan Q, K, dan V dalam satu matrix terkemas; runtime lain memakai matrix terpisah atau sharding tensor-parallel. Invarian arsitekturalnya adalah jumlah KV head hasil proyeksi yang berbeda, bukan nama tensor checkpoint atau urutan packing tertentu.
Tensor parallelism menambahkan batas divisibilitas
Sistem serving sering membagi attention head ke beberapa perangkat. Konfigurasi GQA yang valid secara matematis pada satu perangkat dapat memerlukan penanganan tambahan ketika (H_{kv}) tidak dapat dibagi rata oleh derajat tensor parallelism yang dipilih.
Runtime dapat mereplikasi KV head pada beberapa rank, memakai partisi lain, atau menolak kombinasi jumlah head tertentu. Pilihan tersebut merupakan kontrak runtime. GQA menetapkan pola berbagi antara grup query dan KV head, tetapi tidak menetapkan kebijakan sharding terdistribusi.
Pemisahan yang sama berlaku pada perilaku kernel. Kernel dapat menerima tensor GQA ringkas secara langsung, atau framework dapat memperluas referensi KV agar sesuai dengan antarmuka attention yang lebih seragam. Ukuran cache persisten dan layout eksekusi sementara adalah dua properti berbeda.
Traffic cache tetap bergantung pada implementasi
Jumlah KV head yang lebih sedikit mengurangi data KV berbeda untuk setiap posisi token. Decode path yang memanfaatkan representasi ringkas tersebut dapat menyimpan dan memindahkan lebih sedikit byte. Namun, hal ini tidak menetapkan rasio latensi tertentu.
Latensi decoding juga dipengaruhi bentuk batch, panjang sequence, kernel fusion, hierarki memori, kuantisasi, scheduling, utilisasi perangkat, dan implementasi attention. Runtime yang memperluas data KV bersama sebelum kernel utama dapat memiliki karakteristik traffic berbeda dari runtime yang memproses grouped head secara native.
Batas teknis GQA tetap bersifat struktural: beberapa query head berbagi setiap key-value head. Susunan ini mengurangi jumlah proyeksi KV berbeda dan state KV persisten dibandingkan MHA yang setara pada dimensi lain, sambil mempertahankan beberapa grup KV dan tidak menyatukannya menjadi satu KV head seperti MQA.