GQA · Grouped-Query Attention, chia sẻ đầu KV và ngân sách cache

1Định nghĩa

Grouped-Query Attention (GQA) là thiết kế khối Attention chia toàn bộ H_Q đầu truy vấn (Query heads) thành các nhóm đều nhau, trong đó mỗi nhóm gồm G = H_Q / H_KV đầu Query cùng đọc chung duy nhất 1 cặp đầu Key và Value. Cấu trúc này đứng ở điểm cân bằng giữa hai thái cực: Multi-Head Attention (MHA) (MHA — mỗi đầu Query có riêng 1 đầu Key/Value nên tốn nhiều bộ nhớ nhất) và Multi-Query Attention (MQA) (MQA — tất cả đầu Query ép dùng chung 1 đầu Key/Value duy nhất nên dễ làm giảm độ chính xác).

Khi phục vụ suy luận, GQA cắt giảm cả dung lượng KV cache trên bộ nhớ HBM lẫn lượng byte phải đọc ở mỗi bước decode đúng bằng tỷ lệ gộp nhóm G (ví dụ 32 đầu Query đi với 8 đầu KV giúp giảm 4 lần KV cache). Khác với MLA nén vector KV vào không gian tiềm ẩn, GQA giữ nguyên cấu trúc tensor Key/Value tiêu chuẩn nên dễ dàng kết hợp trực tiếp với Attention cửa sổ trượt (SWA), lượng tử hóa KV cache và chia Tensor Parallelism.

Ví dụ đời thường: GQA giống như cách bố trí tài liệu tra cứu cho 32 chuyên viên ngồi làm việc trong 8 tổ (mỗi tổ 4 người): thay vì in đủ 32 bộ hồ sơ dày cộp cho từng người giữ riêng làm chật kín mặt bàn (như MHA), hoặc chỉ in đúng 1 bộ duy nhất bắt cả 32 người xếp hàng chờ xem chung (như MQA), văn phòng đặt tại mỗi bàn tổ 1 bộ hồ sơ dùng chung cho 4 người trong tổ — giảm ngay 4 lần lượng giấy lưu trữ mà ai cần tra cứu cũng có sẵn ngay trên bàn.

2Luồng

flowchart TD
  step1["Chiếu H_Q đầu Query và H_KV đầu Key/Value"]
  step2["Lưu H_KV đầu vào KV cache trên HBM"]
  step3["Broadcast Key/Value cho nhóm Query trong kernel"]
  step1 --> step2
  step2 --> step3

Luồng mô tả cách GQA giảm kích thước ma trận chiếu K/V, dung lượng KV cache và băng thông đọc trong bước decode.

3Giải pháp

Dung lượng KV cache lý tưởng của GQA cho một chuỗi độ dài T qua L lớp là 2 × L × T × H_KV × d_h × bytes_per_element (chưa tính overhead phân trang và metadata). Khi quy hoạch ngân sách HBM, luôn lấy num_key_value_heads (H_KV) trong config.json thay vì num_attention_heads (H_Q).

Khi triển khai Tensor Parallelism (TP), các đầu attention được chia đều theo chiều head cho các GPU. Nếu số GPU chạy TP nhỏ hơn hoặc bằng H_KV (ví dụ TP = 8 với H_KV = 8), mỗi GPU giữ đúng H_KV / TP đầu KV độc lập; nhưng nếu đẩy TP lớn hơn H_KV (ví dụ TP = 16 khi H_KV = 8), các đầu KV buộc phải nhân bản (replicate) trên nhiều GPU làm giảm hiệu quả tiết kiệm bộ nhớ.

Trong kernel thực thi (như FlashAttention-2/3 hoặc PagedAttention), tuyệt đối không dùng hàm repeat_interleave để nhân bản vật lý tensor Key/Value ra H_Q đầu trên HBM trước khi nhân: kernel chuẩn tự động ánh xạ chỉ số nhóm để đọc 1 đầu KV vào SRAM và dùng chung cho G đầu Query.

4Profiling

• Dung lượng KV cache mỗi token và băng thông đọc decode

  Đối chiếu số byte KV cache thực tế trên HBM với công thức 2 × L × H_KV × d_h × b và đo thời gian bước decode khi tăng cỡ batch/độ dài ngữ cảnh.

• Tỷ lệ nhân bản KV head khi chạy Tensor Parallelism

  Kiểm tra cấu hình TP_size so với num_key_value_heads để phát hiện trường hợp TP_size > H_KV gây nhân bản KV cache giữa các worker.

Mẹo đo lường: ghim số liệu trên workload thật, ghi cả p50/p95/p99 thay vì chỉ trung bình, và luôn ghi lại điều kiện chạy (batch, cache, đồng thời, phần cứng) trước và sau khi tối ưu. Số đo thiếu điều kiện chạy thì không so sánh được.

5Ví dụ

6Thuật ngữ

Grouped-Query Attention (GQA)

Cơ chế chú ý gom G đầu Query dùng chung 1 cặp đầu Key/Value (H_KV = H_Q / G), nằm giữa MHA (H_KV = H_Q) và MQA (H_KV = 1).

KV head broadcast

Trong kernel attention, mỗi cặp Key/Value chỉ đọc 1 lần từ HBM vào SRAM rồi phát chung cho cả nhóm G đầu Query mà không nhân bản dữ liệu trong KV cache.

Mean-pooling uptraining

Kỹ thuật chuyển đổi một checkpoint MHA có sẵn sang GQA bằng cách lấy trung bình các ma trận chiếu Key và Value trong cùng nhóm rồi huấn luyện tiếp một phần nhỏ ngân sách token.

7Bài liên quan

8Nguồn