MQA · Multi-Query Attention, một đầu KV dùng chung và giới hạn Tensor Parallel

1Định nghĩa

Multi-Query Attention (MQA, do Noam Shazeer đề xuất) là biến thể Attention cắt giảm bộ nhớ đệm mạnh tay nhất theo hướng chia sẻ đầu: mô hình vẫn giữ nguyên H_Q đầu truy vấn (Query heads) để đặt nhiều câu hỏi khác nhau, nhưng toàn bộ H_Q đầu đó đều tra cứu trên duy nhất 1 đầu Key và 1 đầu Value dùng chung (H_KV = 1). Có thể xem MQA là trường hợp cực hạn của GQA khi số nhóm bằng 1, đối lập hoàn toàn với MHA nơi mỗi đầu Query có một cặp Key/Value riêng.

Nhờ rút H_KV về 1, MQA giảm dung lượng KV cache và băng thông đọc bộ nhớ HBM ở pha decode đi đúng H_Q lần (ví dụ giảm 32 lần nếu có 32 đầu Query). Tuy nhiên, việc ép mọi đầu Query đọc chung 1 không gian Key/Value duy nhất có thể làm giảm nhẹ chất lượng mô hình lớn và khiến đầu KV duy nhất đó phải nhân bản (replicate) trên mọi GPU khi chạy Tensor Parallelism — lý do khiến nhiều mô hình sau này ưu tiên chọn điểm cân bằng GQA hoặc MLA.

Ví dụ đời thường: MQA giống như việc một phòng ban có 32 chuyên viên nhưng chỉ in đúng 1 cuốn tài liệu tham khảo duy nhất đặt ở bàn giữa phòng cho tất cả 32 người cùng nhìn chung (thay vì 32 cuốn riêng như MHA hay 8 cuốn theo tổ như GQA) — cách này tiết kiệm giấy và diện tích bàn tới mức tối đa, nhưng mọi người chỉ có chung một bản tóm tắt duy nhất để tra cứu.

2Luồng

flowchart TD
  step1["Chiếu H_Q đầu Query và duy nhất 1 đầu Key/Value"]
  step2["Lưu 1 cặp Key/Value duy nhất vào KV cache"]
  step3["Tái sử dụng 1 đầu KV cho mọi đầu Query trên SRAM"]
  step1 --> step2
  step2 --> step3

Luồng mô tả cách MQA dùng duy nhất 1 đầu Key/Value cho toàn bộ H_Q đầu Query nhằm tối thiểu hóa băng thông đọc KV cache.

3Giải pháp

Dung lượng KV cache lý tưởng của MQA chỉ còn 2 × L × T × d_h × bytes_per_element (vì H_KV = 1), nhỏ hơn MHA đúng H_Q lần. Điều này giúp biến bước decode từ chỗ bị nghẽn nặng ở băng thông đọc HBM trở nên nhẹ nhàng hơn rất nhiều ở cỡ batch lớn.

Lưu ý quan trọng khi chạy Tensor Parallelism (TP) trên N GPU: vì chỉ có 1 đầu Key/Value duy nhất (H_KV = 1), engine không thể chia đầu KV theo chiều head cho N GPU mà phải nhân bản (replicate) đầu KV đó trên cả N thiết bị (hoặc phải phân mảnh theo chiều sequence). Do đó trên cụm TP = 8 GPU, dung lượng KV cache trên mỗi GPU của MQA thực tế bằng với GQA có H_KV = 8.

Để bù phần tham số bị giảm ở ma trận chiếu W_K và W_V khi dùng MQA, các kiến trúc thực tế (như PaLM, Falcon hay StarCoder) thường tăng nhẹ chiều rộng của mạng truyền thẳng FFN để giữ nguyên tổng ngân sách tham số.

4Profiling

• Tốc độ decode và băng thông HBM tiết kiệm được

  So sánh thời gian bước decode (TPOT) và lượng byte đọc từ HBM mỗi token giữa MQA (H_KV = 1), GQA và MHA ở cùng độ dài ngữ cảnh.

• Chi phí nhân bản KV cache khi bật Tensor Parallelism

  Kiểm tra bộ nhớ KV cache trên từng GPU khi TP > 1: với MQA, mỗi GPU đều giữ nguyên 1 đầu KV nên tổng KV cache toàn cụm tăng theo TP_size.

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ữ

Multi-Query Attention (MQA)

Cơ chế chú ý giữ H_Q đầu Query nhưng chỉ dùng duy nhất 1 đầu Key và 1 đầu Value (H_KV = 1) cho toàn bộ lớp.

KV head replication under TP

Khi chia Tensor Parallelism trên nhiều GPU, vì MQA chỉ có 1 đầu KV không thể cắt nhỏ theo head, mọi GPU đều phải giữ bản sao của đầu KV duy nhất này.

Arithmetic intensity gain

Nhờ đọc 1 đầu KV từ HBM rồi nhân với cả H_Q đầu Query trong SRAM, MQA tăng mạnh tỷ số phép tính trên mỗi byte bộ nhớ đọc vào.

7Bài liên quan

8Nguồn