Grouped Query Attention (GQA) Là Gì?
Grouped Query Attention (GQA) là kỹ thuật tối ưu attention mechanism trong transformer, chia nhóm các query head cùng chia sẻ chung một bộ Key và Value. GQA nằm giữa Multi-Head Attention (MHA) truyền thống và Multi-Query Attention (MQA), cân bằng giữa chất lượng và tốc độ inference.
Thuật ngữ này trở nên phổ biến khi Meta sử dụng trong Llama 2 70B và Llama 3. Nếu bạn đã nghe Flash Attention giúp tăng tốc attention từ phía thuật toán, thì GQA tối ưu từ phía cấu trúc mô hình.
Trước GQA: Hiểu Multi-Head Attention (MHA)
Trong attention mechanism chuẩn của Transformer, mỗi layer có nhiều “đầu” attention (head). Ví dụ 32 heads. Mỗi head có riêng bộ Query (Q), Key (K), và Value (V) matrices.
Điều này nghĩa là: nếu có 32 heads, bạn cần tính và lưu 32 bộ K và V riêng biệt. Bộ nhớ KV Cache phình to, đặc biệt khi xử lý context dài. Đây là lý do chạy LLM trên GPU tốn nhiều VRAM hơn mức cần thiết.
Vấn Đề GQA Giải Quyết Là Gì?
Năm 2022, nhóm nghiên cứu của Google đề xuất Multi-Query Attention (MQA): tất cả query head dùng chung một bộ K và V duy nhất. Kết quả: bộ nhớ KV Cache giảm mạnh, inference nhanh hơn đáng kể.
Nhưng MQA có vấn đề: chất lượng giảm. Khi tất cả head ép dùng chung một bộ K/V, mô hình mất khả năng chú ý đa chiều. Nghiên cứu cho thấy MQA gây giảm chất lượng rõ rệt so với MHA, đặc biệt ở các bài toán dịch máy và tóm tắt.
GQA đi ra giải quyết bài toán này: thay vì bắt tất cả head dùng chung 1 bộ K/V (như MQA), hoặc mỗi head có riêng (như MHA), GQA chia query head thành nhóm. Mỗi nhóm chia sẻ một bộ K/V.
Ví dụ: 32 query heads chia thành 8 nhóm, mỗi nhóm 4 heads. Thay vì cần 32 bộ K/V (MHA) hay 1 bộ (MQA), bạn chỉ cần 8 bộ. Bộ nhớ giảm 4 lần so với MHA mà chất lượng gần như giữ nguyên.
GQA Hoạt Động Cụ Thể Thế Nào?
Mình giải thích từng bước cho dễ hình dung:
Bước 1: Query vẫn được tính cho tất cả 32 heads bình thường, không thay đổi.
Bước 2: Key và Value chỉ được tính cho 8 nhóm (group heads). Nhóm 1 (gồm head 1-4) dùng chung K1 và V1. Nhóm 2 (head 5-8) dùng chung K2 và V2. Cứ thế cho đến nhóm 8.
Bước 3: Mỗi query head tính attention score với K và V của nhóm nó thuộc về. Head 1-4 đều nhân với K1, V1; head 5-8 nhân với K2, V2.
Kết quả: số phép tính Key/Value giảm 4 lần, bộ nhớ KV Cache giảm 4 lần, nhưng mỗi query head vẫn có “góc nhìn” riêng.
Tại Sao GQA Quan Trọng?
GQA giải quyết bài toán thực tế: chi phí chạy LLM quá đắt. KV Cache chiếm phần lớn VRAM khi inference với context dài. Nếu bạn chạy Llama 2 70B với context 4K tokens, KV Cache alone có thể tốn hơn 10GB VRAM.
Với GQA giảm 4 lần, bạn có thể chạy cùng mô hình trên GPU nhỏ hơn, hoặc tăng context length mà không hết VRAM. Đây là khác biệt giữa “chạy được” và “không chạy được” trên phần cứng cụ thể.
Mình thấy nhiều team Việt Nam muốn deploy LLM self-hosted nhưng bị giới hạn VRAM. GQA chính là một trong những kỹ thuật giúp bridge gap đó.
Những Model Nào Đang Dùng GQA?
GQA đã trở thành standard trong nhiều LLM hiện đại:
Llama 2 70B: Meta dùng GQA với 8 group cho 64 query heads (8 bộ K/V thay vì 64). Đây là lý do Llama 2 70B inference nhanh hơn nhiều so với kích thước.
Llama 3: Tiếp tục dùng GQA, tối ưu hơn nữa cho inference quy mô lớn.
Mistral 7B: Dùng GQA với 4 group, góp phần chạy mượt trên consumer GPU.
Gemma 2: Google cũng áp dụng GQA trong model open-weight của họ.
GQA So Với Các Kỹ Thuật Tối Ưu Khác
GQA không đơn độc. Nó thường kết hợp với các kỹ thuật khác:
GQA + Flash Attention: GQA giảm bộ nhớ K/V, Flash Attention tối ưu cách tính attention. Cùng dùng, hiệu quả cộng hưởng.
GQA + KV Cache tối ưu: GQA giảm kích thước KV Cache trực tiếp, kết hợp với quantization KV Cache (như FP8) để tiết kiệm thêm.
GQA + Pruning: GQA giảm chi phí, pruning giảm số parameter. Cả hai cùng hướng tới mô hình hiệu quả hơn.
Có Nhược Điểm Gì Không?
GQA thực ra là trade-off. Bạn chọn số group càng ít thì tiết kiệm càng nhiều nhưng chất lượng giảm. Nghiên cứu gốc của paper “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints” (2023) cho thấy:
8 groups trở lên: chất lượng gần như MHA. 2-4 groups: chất lượng giảm nhẹ nhưng chấp nhận được. 1 group (tức MQA): chất lượng giảm rõ rệt.
Vấn đề thứ hai: GQA phải được tích hợp vào lúc training. Bạn không thể convert mô hình MHA sang GQA sau khi train xong (hoặc chính xác hơn, cần thêm quá trình continue training để thích nghi). Paper gốc đề xuất phương pháp convert nhưng chất quả không bằng train GQA từ đầu.
Kết Luận
Grouped Query Attention là một trong những tối ưu đơn giản mà hiệu quả nhất trong LLM hiện đại. Không thay đổi kiến trúc cơ bản, chỉ nhóm lại K/V heads, nhưng giảm bộ nhớ và tăng tốc inference đáng kể. Nếu bạn đang tìm hiểu hoặc deploy LLM, GQA là khái niệm đáng để biết.
Cùng với Flash Attention và KV Cache optimization, GQA giúp LLM chạy nhanh hơn, rẻ hơn, tiếp cận được nhiều người hơn. Đó là xu hướng chung: AI không chỉ mạnh hơn mà còn hiệu quả hơn.
