PagedAttention

PagedAttention là thuật toán attention và trình quản lý bộ nhớ dành cho việc phục vụ mô hình ngôn ngữ lớn (LLM); thuật toán này lưu KV cache của từng request trong các block kích thước cố định nhỏ, đặt tại bất kỳ vị trí nào trong bộ nhớ GPU, rồi theo dõi chúng bằng bảng block cho mỗi request, tương tự cách hệ điều hành ánh xạ các trang bộ nhớ ảo. Woosuk Kwon, Zhuohan Li và các đồng nghiệp tại UC Berkeley cùng các trường đại học đối tác giới thiệu thuật toán này cùng với engine phục vụ vLLM trong bài báo Efficient Memory Management for Large Language Model Serving with PagedAttention, được trình bày tại SOSP 2023. Loại bỏ phân mảnh do cấp phát cache liên tục giúp một GPU chứa đồng thời nhiều chuỗi hơn, và bài báo ghi nhận thông lượng cao hơn 2–4 lần so với FasterTransformer và Orca ở cùng mức độ trễ.

Vấn đề bộ nhớ mà PagedAttention giải quyết#

Trong quá trình tạo văn bản, một transformer giữ lại vector khóa và giá trị của mọi token trước đó để không phải tính lại chúng. Cache này tăng thêm một mục cho mỗi token ở mỗi layer; một chuỗi đơn của model LLaMA 13 tỷ tham số có thể chiếm đến 1.7 GB. Độ dài đầu ra chưa xác định khi request đến, vì vậy các hệ thống phục vụ đời trước dành riêng cho mỗi request một vùng liên tục, được cấp phát theo độ dài tối đa có thể. Phần lớn vùng đó không được sử dụng, còn các khoảng trống giữa những vùng quá nhỏ để tái sử dụng. Nhóm vLLM đo được rằng các hệ thống như vậy lãng phí 60–80% bộ nhớ KV cache, khiến batch size bị giới hạn từ lâu trước khi hết compute.

Cách phân trang hoạt động#

PagedAttention chia cache thành các block, mỗi block chứa khóa và giá trị của một số lượng token cố định; trong vLLM, mặc định là 16 token. Request chỉ nhận block mới khi block cuối đã đầy, từ bất kỳ vị trí nào có block trống; bảng block của request ghi lại thứ tự. Kernel cơ chế attention dựa vào bảng đó để tập hợp khóa và giá trị theo từng block, vì vậy cache của một chuỗi không cần nằm liên tục trong bộ nhớ. Lãng phí chỉ giới hạn ở phần chưa lấp đầy tại cuối block cuối cùng của mỗi chuỗi, dưới 4% theo số đo của vLLM.

Đoạn mã minh họa PyTorch này lưu khóa của một request trong các block rải rác thuộc một pool dùng chung, rồi đọc lại thông qua bảng block của request:

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

Kết quả in ra bảng block [41, 7, 19] và keys intact: True, rồi cho thấy request dùng phân trang chỉ có 8 vị trí chưa sử dụng, so với 2,008 vị trí trong vùng dành riêng cho 2,048 token.

Chia sẻ block giữa các chuỗi#

Bảng block cũng cho phép nhiều chuỗi trỏ đến cùng một block vật lý. Lấy mẫu song song, beam search và các request bắt đầu bằng cùng một system prompt có thể dùng lại một bản sao của tiền tố chung, với số lượng tham chiếu được lưu cho mỗi block. Khi một chuỗi cần ghi vào block vẫn đang được chuỗi khác sử dụng, engine sẽ sao chép block đó trước; đây là bước copy-on-write mượn từ hệ điều hành.

Các tác giả vLLM báo cáo rằng cách chia sẻ này giảm mức sử dụng bộ nhớ của lấy mẫu song song và beam search đến 55%, tương đương thông lượng cao hơn đến 2.2 lần. Hash toàn bộ block mở rộng ý tưởng này qua các request, đó là cách vLLM triển khai cache prompt.

Các kỹ thuật liên quan#

PagedAttention quyết định vị trí lưu khóa và giá trị đã cache, đồng thời kết hợp với các phương pháp thay đổi lượng dữ liệu cần lưu hoặc tốc độ đọc dữ liệu:

• Flash Attention: Chia nhỏ phép tính attention để toàn bộ ma trận điểm số không bao giờ được ghi vào bộ nhớ chính của GPU. Các engine hiện nay chạy kernel kiểu FlashAttention để đọc dữ liệu từ những block phân trang.
• Batching liên tục: Tiếp nhận request mới ngay khi một request khác hoàn tất; phân trang lập tức trả các block của chuỗi đã hoàn tất về pool trống, nhờ đó có chỗ cho request mới.
• Attention truy vấn theo nhóm: Chia sẻ các head khóa và giá trị giữa các head truy vấn, làm giảm kích thước nội dung của mỗi block.
• Lượng tử hóa cache FP8: Lưu khóa và giá trị ở độ chính xác 8 bit, giảm gần một nửa bộ nhớ cache so với định dạng 16 bit.

Giao diện LLM của Ultralytics chuyển tiếp nguyên trạng các đối số request, vì vậy một lệnh gọi Chat Completions có n=3 được gửi qua base_url đến máy chủ vLLM cục bộ sẽ trả về ba mẫu dùng chung các block đã cache của prompt; chỉ những phần tiếp nối khác nhau mới cần bộ nhớ mới.