Grouped Query Attention (GQA)

Grouped Query Attention (GQA) là một thiết kế attention trong đó nhiều query head chia sẻ một tập hợp nhỏ hơn gồm các key head và value head. Thiết kế này duy trì các góc nhìn đa dạng của query đa đầu trong khi giảm bộ nhớ và dữ liệu di chuyển cần thiết cho key và value. GQA đặc biệt có giá trị trong quá trình suy luận tự hồi quy (autoregressive inference), nơi model tạo ra từng token một và đọc lặp lại các trạng thái attention được lưu trữ.

Cách GQA Hoạt Động#

Trong một attention mechanism, mỗi token được chiếu thành ba biểu diễn:

• Query: Mô tả thông tin mà token hiện tại đang tìm kiếm.
• Key: Mô tả những gì mỗi token có sẵn đại diện.
• Value: Chứa thông tin được truy xuất khi một query khớp với một key.

Trong self-attention thông thường, các phép chiếu này được chia thành các head để model có thể học các mối quan hệ khác nhau song song. GQA giữ lại nhiều query head nhưng gán các nhóm head này cho các key head và value head được chia sẻ. Ví dụ, một layer attention có thể sử dụng tám query head và hai key-value head. Khi đó, mỗi key-value head sẽ phục vụ bốn query head.

Phép tính attention bản thân nó vẫn là scaled dot-product attention. Điều thay đổi là sự sắp xếp head và lượng dữ liệu key-value được tạo ra, lưu trữ và đọc. Do đó, các framework như PyTorch scaled dot-product attention yêu cầu số lượng query head phải chia hết cho số lượng key-value head khi GQA được bật. (docs.pytorch.org)

GQA so với Multi-Head và Multi-Query Attention#

GQA nằm ở giữa multi-head attention và multi-query attention:

• Multi-head attention: Mỗi query head có key head và value head riêng. Điều này mang lại tính độc lập tối đa cho head nhưng tạo ra yêu cầu bộ nhớ key-value lớn nhất.
• Grouped query attention: Một số query head chia sẻ chung một key-value head. Phương pháp này cân bằng giữa năng lực biểu diễn và hiệu quả bộ nhớ.
• Multi-query attention: Tất cả query head chia sẻ một key head và một value head. Điều này giảm thiểu việc lưu trữ key-value nhưng cung cấp độ đa dạng key-value thấp hơn.

Nếu một layer có 32 query head, multi-head attention có thể sử dụng 32 key-value head, GQA có thể sử dụng 8 và multi-query attention sử dụng 1. GQA không phải là sự thay thế cho Transformer architecture tổng thể; nó là một cấu hình khả dĩ bên trong một attention layer của Transformer. Tài liệu multi-head, multi-query, và grouped-query attention documentation của NVIDIA mô tả các biến thể này khác nhau chủ yếu ở số lượng key-value head phục vụ cho các query head. (nvidia.github.io)

Tại sao GQA Lại Quan Trọng#

Trong quá trình tạo tự hồi quy, model lưu trữ các key và value trước đó trong một KV cache. Với multi-head attention tiêu chuẩn, mọi attention head đều đóng góp các key và value được cache riêng biệt. Do đó, các chuỗi dài, batch lớn và nhiều layer model có thể tiêu thụ một lượng đáng kể bộ nhớ GPU.

Vì GQA sử dụng ít key-value head hơn, bộ nhớ cache của nó nhỏ hơn tỷ lệ thuận với việc giảm số lượng head đó. Điều này có thể cải thiện:

• Dung lượng bộ nhớ: Các prompt dài hơn hoặc nhiều yêu cầu đồng thời hơn có thể vừa vặn trong bộ nhớ khả dụng.
• Thông lượng tạo (Generation throughput): Cần ít dữ liệu key-value hơn được đọc cho mỗi token được tạo ra.
• Inference latency: Lưu lượng truy cập bộ nhớ giảm có thể rút ngắn thời gian xử lý giữa các token.
• Khả năng mở rộng context window: Việc phục vụ các chuỗi dài trở nên khả thi hơn.

Những lợi ích này phụ thuộc vào phần cứng, độ dài chuỗi, kích thước batch và hỗ trợ kernel. Các runtime production vẫn cần phân bổ cache hiệu quả; ví dụ, hệ thống TensorRT-LLM KV cache system của NVIDIA kết hợp hỗ trợ GQA với caching dựa trên block, tái sử dụng và offload. (nvidia.github.io)

Các ứng dụng trong thực tế#

Trợ lý ngữ cảnh dài (Long-context assistants): Trợ lý lập trình có thể xử lý hàng nghìn token mã nguồn trước khi tạo câu trả lời. GQA làm giảm trạng thái key-value được cache liên quan đến lịch sử đó, cho phép server xử lý các file dài hơn hoặc nhiều người dùng đồng thời hơn mà không làm tăng tỷ lệ bộ nhớ GPU.

Trợ lý hình ảnh và video đa phương thức: Một vision-language model có thể biểu diễn các bản vá hình ảnh (image patches) hoặc khung hình video dưới dạng chuỗi token dài. GQA làm giảm áp lực bộ nhớ cache attention sau khi các token trực quan này đi vào bộ giải mã ngôn ngữ (language decoder), có khả năng để lại nhiều bộ nhớ hơn cho các hình ảnh, khung hình hoặc yêu cầu bổ sung. Điều này khác với Vision Transformer tiêu chuẩn, nơi attention có thể xử lý tất cả các bản vá hình ảnh song song mà không cần KV caching tự hồi quy.

Triển Khai Thực Tế và Các Đánh Đổi#

PyTorch cung cấp GQA thông qua enable_gqa=True. Ví dụ CUDA được tài liệu hóa này sử dụng tám query head và hai key-value head được chia sẻ:

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

assert torch.cuda.is_available(), "A CUDA device is required."

query = torch.randn(1, 8, 16, 32, device="cuda")
key = torch.randn(1, 2, 16, 32, device="cuda")
value = torch.randn(1, 2, 16, 32, device="cuda")

# Four query heads share each key-value head.
with sdpa_kernel(SDPBackend.MATH):
    output = F.scaled_dot_product_attention(
        query,
        key,
        value,
        enable_gqa=True,
    )

print(output.shape)

Kết quả đầu ra giữ lại tám query head, trong khi key và value chỉ sử dụng hai head. Trình chọn backend attention PyTorch attention backend selector có thể kiểm soát kernel được hỗ trợ nào thực hiện thao tác, trong khi hướng dẫn khối xây dựng PyTorch Transformer building-block guide cung cấp bối cảnh triển khai rộng hơn. (docs.pytorch.org)

GQA phải được chọn khi thiết kế hoặc điều chỉnh model; đây thường không phải là một công tắc trong thời điểm suy luận cho một checkpoint có sẵn. Các lập trình viên nên xác minh sự hỗ trợ của framework, tính chia hết của head, hành vi số học, mức sử dụng bộ nhớ và chất lượng tác vụ. Đối với việc triển khai thị giác máy tính, các quy trình bổ sung như Ultralytics TensorRT integration và Ultralytics benchmark mode giúp đo lường xem các tối ưu hóa kiến trúc và runtime có mang lại những cải thiện đáng kể trên phần cứng mục tiêu hay không.