PagedAttention Là Gì?
PagedAttention là thuật toán quản lý bộ nhớ KV cache cho mô hình ngôn ngữ lớn, lấy cảm hứng từ cơ chế phân trang (paging) của hệ điều hành. Thay vì cấp phát một vùng nhớ liên tục cho từng yêu cầu, PagedAttention chia KV cache thành các block nhỏ và cấp phát theo nhu cầu, giúp giảm thiểu lãng phí bộ nhớ GPU xuống chỉ còn khoảng 4%.
Thuật toán này được giới thiệu trong bài báo “Efficient Memory Management for Large Language Model Serving with PagedAttention” tại hội nghị SOSP 2023, do nhóm nghiên cứu tại UC Berkeley phát triển. Nó là nền tảng cốt lõi của vLLM, một trong những hệ thống serve LLM mã nguồn mở phổ biến nhất hiện nay.
Vấn Đề PagedAttention Giải Quyết Là Gì?
Khi một LLM sinh văn bản, nó cần lưu trữ key và value của mọi token đã xử lý vào cái gọi là KV cache. Vấn đề là cách truyền thống cấp phát bộ nhớ cho KV cache rất wasteful: hệ thống reserve trước một vùng nhớ liên tục đủ lớn cho độ dài tối đa mà yêu cầu có thể đạt tới.
Nhưng thực tế, phần lớn yêu cầu kết thúc sớm hơn nhiều so với max length. Ví dụ: bạn reserve buffer cho 2.048 token, nhưng câu trả lời chỉ dài 100 token. Kết quả là 1.948 slot bị reserve mà không ai dùng được. Nghiên cứu cho thấy 60-80% bộ nhớ KV cache bị lãng phí theo cách này.
Đây gọi là internal fragmentation. Thêm vào đó, khi các yêu cầu vào ra liên tục, các khoảng trống giữa chúng tạo ra external fragmentation. Tổng lại, GPU có thể đang ngồi trên đống bộ nhớ trống nhưng vẫn phải từ chối yêu cầu mới.
PagedAttention Hoạt Động Như Thế Nào?
Ý tưởng rất elegant: mượn y nguyên cơ chế virtual memory paging mà hệ điều hành dùng cho RAM từ 50 năm nay. Thay vì một buffer liên tục, KV cache được chia thành các block cố định, ví dụ mỗi block chứa 16 token.
Mỗi yêu cầu duy trì một block table, ánh xạ từ vị trí logic sang physical block trong GPU memory. Giống hệt cách page table của OS maps virtual page sang physical frame. Các block có thể nằm rải rác bất kỳ đâu trong memory, không cần liên tục.
Kernel attention được viết lại để gather keys và values thông qua block table thay vì giả định chúng nằm liên tục. Chi phí lookup thêm rất nhỏ so với lượng bộ nhớ tiết kiệm được.
Lợi Ích Chính Của PagedAttention Là Gì?
1. Giảm lãng phí bộ nhớ: Thay vì lãng phí cả ngàn slot, mỗi yêu cầu chỉ lãng phí tối đa một block (khoảng 15 slot với block size 16). Nghiên cứu cho thấy memory waste giảm từ 60-80% xuống chỉ còn khoảng 4%.
2. Tăng throughput 2-4 lần: Bộ nhớ tiết kiệm được biến thành batch size lớn hơn. Nhiều yêu cầu hơn chạy song song trên cùng GPU, trực tiếp tăng tokens/giây.
3. Copy-on-write sharing: Nhờ block table, nhiều sequence có thể chia sẻ cùng physical block cho prefix chung. Ví dụ: parallel sampling tạo 4 câu trả lời từ 1 prompt, prompt đó chỉ lưu 1 lần. Khi một sequence diverge, chỉ block bị thay đổi mới được copy.
PagedAttention Khác Gì So Với Cách Truyền Thống?
Cách truyền thống giống như đặt trước một phòng họp lớn cho mỗi cuộc họp, dù cuộc họp chỉ có 3 người. PagedAttention giống như chia phòng thành các booth nhỏ, cấp phát booth theo số người thực tế. Phòng trống ngay lập tức được tái sử dụng.
Điểm hay nhất là không cần thay đổi gì model. PagedAttention là optimization ở tầng serving, model weights và architecture giữ nguyên. Kiến trúc này hoạt động với bất kỳ transformer model nào.
Ai Nên Dùng PagedAttention?
Nếu bạn đang deploy LLM production, đặc biệt với workload có độ dài sequence biến động mạnh, PagedAttention gần như là must-have. vLLM, hệ thống serve được xây dựng trên PagedAttention, đã trở thành standard de facto cho LLM serving mã nguồn mở.
Các use case được hưởng lợi nhiều nhất:
- API chatbot với conversation dài ngắn khác nhau
- Parallel sampling và beam search (hưởng lợi từ prefix sharing)
- Multi-tenant serving với system prompt chung
- Batch processing với độ dài input thay đổi
PagedAttention Có Nhược Điểm Gì Không?
Có, nhưng rất nhỏ. Việc gather keys/values qua block table thêm một chút overhead so với đọc memory liên tục. Ở batch size 1 (chỉ 1 yêu cầu), chi phí này có thể thấy được. Nhưng ở workload thực tế với batch size lớn, lợi ích bộ nhớ tiết kiệm được áp đảo hoàn toàn.
Ngoài ra, PagedAttention yêu cầu custom CUDA kernel để hoạt động hiệu quả, nên không phải plug-and-play trên mọi hardware. Nhưng vLLM đã handle phần này, nên nếu dùng vLLM thì bạn không cần lo.
Kết Luận
PagedAttention là một trong những optimization đẹp nhất mà mình thấy trong không gian LLM serving. Nó lấy một ý tưởng cũ (OS paging), áp dụng vào bài toán mới (KV cache management), và giải quyết chính bottleneck của inference. Kết quả: throughput tăng 2-4 lần, memory waste giảm từ 80% xuống 4%, không cần thay đổi model.
Nếu bạn đang tự host LLM, dùng vLLM là cách dễ nhất để được hưởng lợi từ PagedAttention. Nó open source, widely adopted, và đã chứng minh trong production ở quy mô lớn.
