MHA · Multi-Head Attention, độc lập đầu Q/K/V và chi phí KV cache

1Định nghĩa

Multi-Head Attention (MHA) là cấu trúc Attention nguyên bản của kiến trúc Transformer: thay vì thực hiện một phép chú ý duy nhất trên toàn bộ vector ẩn chiều d_model, mô hình chia vector thành H đầu (heads) chạy song song với chiều mỗi đầu d_h = d_model / H. Trong MHA, mỗi đầu sở hữu riêng một bộ ba ma trận chiếu Query, Key và Value hoàn toàn độc lập (H_Q = H_KV = H), giúp mô hình đồng thời theo dõi nhiều mối quan hệ ngữ pháp và ngữ nghĩa khác nhau trên cùng một chuỗi.

Tuy mang lại khả năng biểu diễn dồi dào khi huấn luyện, việc giữ đủ H đầu Key và H đầu Value riêng biệt khiến MHA trở thành điểm nghẽn bộ nhớ lớn nhất khi phục vụ suy luận: dung lượng KV cache trên bộ nhớ HBM và lượng dữ liệu phải đọc ở mỗi bước decode đều tăng tỷ lệ thuận với tổng số đầu H. Chính chi phí băng thông đọc KV cache này là lý do hầu hết các LLM lớn hiện nay chuyển từ MHA sang GQA, MQA hoặc MLA.

Ví dụ đời thường: MHA giống như một đoàn thanh tra gồm 32 chuyên viên cùng đọc lại một tập hồ sơ vụ việc, trong đó mỗi chuyên viên tự mang theo một cuốn sổ ghi chép (Key/Value) dày cộp của riêng mình để đánh dấu góc nhìn chuyên môn riêng — cách làm này rất kỹ lưỡng và đa chiều, nhưng bàn làm việc phải đủ rộng để chứa đủ cả 32 cuốn sổ riêng biệt đó.

2Luồng

flowchart TD
  step1["Chiếu thành H bộ Query, Key, Value riêng"]
  step2["Ghi đủ H đầu Key/Value vào KV cache"]
  step3["Tính attention song song và ghép qua W_O"]
  step1 --> step2
  step2 --> step3

Luồng mô tả cách MHA tách H đầu độc lập, lưu toàn bộ H cặp Key/Value vào cache và ghép kết quả ở đầu ra.

3Giải pháp

Với mô hình dùng MHA, mỗi token ở mỗi lớp cần lưu 2 × H × d_h phần tử vào KV cache. Khi phục vụ chuỗi dài hoặc batch lớn, phần bộ nhớ này nhanh chóng vượt qua cả dung lượng trọng số mô hình trên HBM.

Khi chia mô hình qua nhiều GPU bằng Tensor Parallelism, MHA rất dễ chia đều vì cả Query, Key và Value đều có cùng số đầu H: chỉ cần cắt ma trận W_Q, W_K, W_V theo cột và cắt W_O theo hàng cho N thiết bị (miễn H chia hết cho N), sau đó gọi một phép All-Reduce duy nhất sau W_O.

Nếu đang vận hành một checkpoint MHA cũ bị nghẽn ở băng thông đọc KV cache lúc decode, hãy ưu tiên bật lượng tử hóa KV cache (FP8/INT8), dùng FlashAttention / PagedAttention hoặc uptrain sang GQA.

4Profiling

• Tỷ trọng KV cache so với trọng số trên HBM

  Tính 2 × L × T × H × d_h × b theo độ dài chuỗi T và cỡ batch để xác định ngưỡng bắt đầu tràn HBM khi dùng MHA.

• Băng thông đọc HBM ở bước decode

  Đo thời gian kernel attention khi decode chuỗi dài; với MHA, lượng byte đọc KV cache trên mỗi đầu Query cao gấp G lần so với GQA.

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-Head Attention (MHA)

Cơ chế chú ý nhiều đầu tiêu chuẩn trong đó số đầu Key và Value bằng đúng số đầu Query (H_KV = H_Q = H).

Head dimension (d_h)

Số chiều của mỗi đầu attention (thường là 64 hoặc 128), sao cho H × d_h bằng chiều rộng chiếu của lớp attention.

Output projection (W_O)

Ma trận chiếu cuối cùng ghép (concat) kết quả của cả H đầu attention lại và trộn chúng về lại không gian vector ẩn d_model.

7Bài liên quan

8Nguồn