PagedAttention (vLLM) — Virtual memory cho AI, waste giảm từ 70% → 4%

Giải thích PagedAttention trong vLLM: cách áp dụng nguyên lý bộ nhớ ảo của OS để giảm lãng phí KV cache từ 70% xuống 4%, tăng throughput 24x.

Chạy LLM 70B trên GPU cho hàng nghìn user đồng thời là bài toán kinh tế. Chi phí GPU không phải nằm ở FLOPs, mà ở VRAM — và cách chúng ta cấp phát bộ nhớ truyền thống đang lãng phí đến 70% dung lượng. PagedAttention mang nguyên lý virtual memory của hệ điều hành vào inference: thay vì đặt chỗ trước một dải bộ nhớ liên tục, nó cấp phát từng "trang" (page) nhỏ khi cần, giảm waste xuống dưới 4% và tăng throughput serving lên 24 lần.

Vấn đề

Cái bẫy của pre-allocation.

Truyền thống, khi một request đến, hệ thống inference sẽ cấp phát một khối VRAM liên tục cho toàn bộ max_seq_len (ví dụ 4096 tokens) ngay từ đầu. Nhưng thực tế, đa số câu trả lời chỉ dài 50-200 tokens. Phần còn lại — hàng nghìn slot KV cache — cứ nằm đó chiếm chỗ như ghế trống trong rạp chiếu phim đã đặt vé nhưng không ai ngồi. Đây gọi là internal fragmentation.

Vấn đề thứ hai là external fragmentation: khi nhiều request kết thúc ở độ dài khác nhau, chúng để lại những "lỗ hổng" nhỏ trong bộ nhớ GPU. Tổng dung lượng trống có thể đủ cho request mới, nhưng vì không liên tục, hệ thống từ chối request — dù VRAM còn thừa. Kết quả: GPU "đói" bộ nhớ trong khi hàng GB VRAM đang bị cắt nhỏ vô dụng.

Kết hợp lại, các hệ thống inference truyền thống (HF Transformers, TGI nguyên bản) lãng phí 60-80% VRAM chỉ cho việc đặt chỗ thừa. Đây là lý do bạn không thể chạy batch lớn dù model chỉ chiếm một nửa GPU.

Ý tưởng cốt lõi

"Attention không quan tâm địa chỉ vật lý."

PagedAttention bắt đầu từ một câu hỏi đơn giản: Tại sao KV cache phải nằm liền nhau trong VRAM? Khi tính attention, model chỉ cần biết "vector K của token số 5 ở đâu", chứ không cần biết nó có nằm ngay cạnh vector K của token số 6 hay không. Miễn là kernel CUDA tìm được nó, thứ tự vật lý là không quan trọng.

Ẩn dụ "Bãi đỗ xe ảo":

Tưởng tượng bãi đỗ xe truyền thống là một dải liên tục: khi xe vào, bạn chặn một chỗ dài bằng cả tuần dù xe chỉ đỗ 2 tiếng. PagedAttention như hệ thống vé điện tử hiện đại: bãi xe được chia thành từng chỗ nhỏ (block), xe nào đến được phát chỗ đó, ra khỏi thì thu hồi. Xe đi nối tiếp nhau không cần đỗ cạnh nhau — chỉ cần có bản đồ (page table) ghi rõ "xe A ở chỗ 5, chỗ 7, chỗ 12".

Cơ chế Block Table:

Block granularity: KV cache được chia thành các block cố định (thường 16 tokens). Mỗi block là một "trang" bộ nhớ.
Logical → Physical mapping: Mỗi sequence có một bảng tra (block table) map vị trí logic (token 0-15, 16-31...) sang địa chỉ vật lý trong VRAM. Hai token liên tiếp logic có thể nằm ở hai góc khác nhau của GPU.
On-demand allocation: Block chỉ được cấp phát khi token thực sự được sinh ra. Request dự đoán 100 tokens chỉ chiếm 7 block (16×7=112), thay vì 256 block cho 4096 tokens tối đa.

Điểm "aha" nằm ở đây: chúng ta đã tách rời không gian địa chỉ logic khỏi bố trí vật lý, giống như OS ảo hóa bộ nhớ RAM cho process. Điều này cho phép dùng lại ngay cả những mảnh bộ nhớ nhỏ (các block giải phóng giữa chừng) mà không cần defragmentation.

Copy-on-Write cho Beam Search:

Một tính năng tinh tế khác: khi dùng beam search hoặc parallel sampling (sinh nhiều output từ một prompt), các nhánh chia sẻ chung prefix (các token đầu tiên giống nhau). PagedAttention cho phép các sequence khác nhau chia sẻ vật lý cùng một block cho phần prefix, chỉ tách ra (copy-on-write) khi token bắt đầu khác nhau. Điều này tiết kiệm thêm bộ nhớ cho multi-sample decoding.

Tại sao nó hoạt động

Kernel scatter-gather: Để attention hoạt động trên bộ nhớ không liên tục, vLLM viết các CUDA kernel tùy chỉnh thay vì dùng thư viện chuẩn. Kernel này nhận danh sách con trỏ (block table) và "thu thập" (gather) các vector K/V từ các địa chỉ rải rác vào shared memory trước khi tính softmax. Overhead của việc nhảy vùng nhớ nhỏ hơn nhiều so với lợi ích của việc không bỏ trống VRAM.

Tích hợp Continuous Batching: PagedAttention là nền tảng cho Continuous Batching. Khi request mới đến giữa chừng, hệ thống chỉ cần cấp thêm block trống rải rác, thay vì tìm một dải liên tục lớn. Sự kết hợp này cho phép GPU luôn làm việc ở công suất tối đa, xử lý request mới mà không đợi cũ xong.

Block size tuning: Kích thước block (thường 16) là trade-off. Quá nhỏ → overhead quản lý block table tăng; quá lớn → internal fragmentation quay lại (waste trong block cuối). 16 là "sweet spot" thực nghiệm cho transformer hiện đại.

Ý nghĩa thực tế

Hiệu suất thực tế:

Giảm waste: Từ 60-80% xuống dưới 4% (chỉ còn lãng phí trong block cuối cùng của mỗi sequence).
Throughput: LMSYS Chatbot Arena báo cáo giảm 50% số GPU cần thiết để phục vụ cùng lượng request, tương đương tăng throughput 2-3x so với TGI và 24x so với HF Transformers nguyên bản.
Chi phí: Cho phép chạy model lớn (70B) trên ít GPU hơn, hoặc phục vụ batch size lớn hơn trên cùng phần cứng.

Giới hạn cần hiểu: PagedAttention chỉ giải quyết memory capacity, không giải quyết memory bandwidth hay compute FLOPs. Attention vẫn là O(n²) về số phép tính — nó chỉ cho phép bạn nhét nhiều sequence hơn vào GPU cùng lúc. Nếu bạn đã bị giới hạn bởi FLOPs (rất hiếm với LLM hiện đại), PagedAttention không giúp nhanh hơn.

Block size cũng là hyperparameter: với context rất ngắn (< 1K), overhead quản lý block có thể không đáng; với context dài (100K+), block size 16 tạo ra quá nhiều block cần tra cứu.

Đào sâu hơn

Paper gốc: Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al., 2023. Giới thiệu block table và copy-on-write semantics cho KV cache.

Bài liên quan TroiSinh:

Cùng cụm (KV Cache & Inference):

KV Cache — Hiểu sâu về cơ chế cache K/V và tại sao nó lại là bottleneck bộ nhớ
Continuous Batching — Cách PagedAttention kết hợp với batching động để GPU không bao giờ "nghỉ"
Prefix Caching — Tối ưu thêm cho system prompt dùng lại nhiều lần
Speculative Decoding — Kết hợp PagedAttention với draft model để giảm latency đầu ra
Beam Search — Tận dụng copy-on-write của PagedAttention cho beam search hiệu quả
Constrained Decoding — Ép output theo grammar trong khi vẫn tận dụng block-based memory

Đọc tiếp (Hiệu quả Attention & Nén):

Flash Attention — Giải quyết O(n²) compute và bandwidth thay vì chỉ memory capacity
Quantization (AWQ/GPTQ) — Giảm memory footprint của weight (thay vì KV cache) để chạy model lớn hơn
Inference Frontier — Các kỹ thuật inference cấp tiến (Level 2) khi PagedAttention không còn đủ

Mở rộng:

vLLM Documentation — Chi tiết kỹ thuật kernel implementation (lưu ý: tài liệu này ghi chép lịch sử thiết kế, hiện tại vLLM dùng optimized kernels trong csrc/attention/)
Blog vLLM — Giải thích từ đội ngũ phát triển về ứng dụng thực tế tại LMSYS Chatbot Arena