Khám phá vLLM: Giải phẫu hệ thống suy luận LLM thông lượng cao

Khám phá vLLM: Giải phẫu hệ thống suy luận LLM thông lượng cao: LLM Engine & Engine Core

Trong loạt bài viết này, chúng ta sẽ đi tìm hiểu tất cả các thành phần cốt lõi và các tính năng nâng cao tạo nên một hệ thống suy luận LLM hiện đại có thông lượng cao. Cụ thể, chúng ta sẽ phân tích cách vLLM hoạt động.

Bao gồm 5 phần:

LLM engine & engine core: nền tảng của vLLM (lên lịch, paged attention, continuous batching…)
Các tính năng nâng cao: điền trước theo từng khối (chunked prefill), bộ nhớ đệm tiền tố (prefix caching), giải mã có hướng dẫn (guided decoding), giải mã dự đoán (speculative decoding), P/D phân tách (disaggregated P/D)
Scaling up: từ một GPU đến nhiều GPU
Serving layer: lớp phục vụ phân tán / đồng thời
Benchmarks and auto-tuning: Điểm chuẩn và tự động tinh chỉnh

Bài này chúng ta cùng tìm hiểu phần thứ nhất: LLM Engine & Engine Core

LLM engine là khối xây dựng cơ bản nhất của vLLM. Tự bản thân nó đã cho phép suy luận với thông lượng cao, nhưng mới chỉ trong bối cảnh offline. Bạn chưa thể dùng nó để phục vụ người dùng qua web.

Ví dụ đơn giản:

Cấu hình này là:

offline (không có web hay hệ thống phân tán),
đồng bộ (tất cả đều chạy trong một tiến trình chặn),
một GPU (no data/model/pipeline/expert parallelism; DP/TP/PP/EP = 1)
dùng transformer chuẩn.

Từ đây, bài viết sẽ dần xây dựng lên thành một hệ thống online, bất đồng bộ, nhiều GPU, nhiều node.

Trong ví dụ này, chúng ta thực hiện hai việc:

Khởi tạo một engine
Gọi phương thức “generate” trên engine đó để lấy mẫu từ các câu lệnh gợi ý (prompt) được cung cấp

1. Hàm khởi tạo LLM Engine

Các thành phần chính của engine:

vLLM config: chứa mọi tham số cấu hình mô hình, cache, song song hóa…
Processor: biến đầu vào thô thành EngineCoreRequest (kiểm tra, tokenize, xử lý).
Engine core client: trong ví dụ này là InprocClient (gần như tương đương EngineCore).
Output processor: biến EngineCoreOutput thành kết quả người dùng nhìn thấy.

Bên trong Engine Core lại gồm:

Model Executor: điều khiển các forward pass trên mô hình.
Structured Output Manager: dùng cho guided decoding.
Scheduler: quyết định request nào được chạy ở bước tiếp theo.

Scheduler chứa:

Chính sách FCFS hoặc priority.
Hàng đợi waiting và running.
KV cache manager: trái tim của paged attention.

KV cache manager duy trì một free_block_queue – một hồ chứa các block KV cache còn trống (có thể lên tới hàng trăm nghìn block tùy VRAM và kích thước block).

Các thành phần cốt lõi được mô tả trong phần này và mối quan hệ giữa chúng.

lưu ý: Kích thước khối cho một lớp transformer tiêu chuẩn (không phải MLA [4]) được tính như sau: 2 (key/value) * block_size (mặc định=16) * num_kv_heads * head_size * dtype_num_bytes (ví dụ: 2 đối với bf16)

Trong quá trình xây dựng model executor, một đối tượng Worker được tạo ra và ba thủ tục quan trọng được thực hiện. (Sau này, khi dùng MultiProcExecutor, chính ba thủ tục này sẽ được chạy độc lập trên từng tiến trình worker ở các GPU khác nhau.)

1.1. Khởi tạo thiết bị (Init device)

Gán một thiết bị CUDA (ví dụ “cuda:0”) cho worker và kiểm tra kiểu dữ liệu (dtype) của mô hình có được hỗ trợ hay không (ví dụ bf16).
Xác minh có đủ VRAM hay không, dựa trên giá trị gpu_memory_utilization được yêu cầu (ví dụ 0.8 → 80% tổng VRAM).
Thiết lập các cấu hình phân tán (DP / TP / PP / EP, v.v.).
Khởi tạo một đối tượng model_runner (chứa sampler, KV cache và các bộ đệm cho forward pass như input_ids, positions, v.v.).
Khởi tạo một đối tượng InputBatch (chứa các bộ đệm phía CPU cho forward pass, các block table dùng để đánh chỉ mục KV cache, metadata của sampling, v.v.).

1.2. Nạp mô hình (Load model)

Khởi tạo kiến trúc mô hình.
Nạp trọng số của mô hình.
Gọi model.eval() (chế độ suy luận của PyTorch).
Tùy chọn: gọi torch.compile() cho mô hình.

1.3. Khởi tạo KV cache

Lấy đặc tả KV cache cho từng layer. Trước đây, nó luôn là FullAttentionSpec (transformer đồng nhất), nhưng với các mô hình lai (sliding window, Transformer/SSM như Jamba) thì việc này trở nên phức tạp hơn (xem Jenga [5]).
Chạy một forward pass giả / phục vụ profiling và chụp snapshot bộ nhớ GPU để tính xem có bao nhiêu block KV cache có thể chứa vừa trong VRAM khả dụng.
Cấp phát, reshape và gắn các tensor KV cache vào các layer attention.
Chuẩn bị metadata attention (ví dụ đặt backend thành FlashAttention) để sau này các kernel sử dụng trong forward pass.
Trừ khi truyền cờ –enforce-eager, với mỗi kích thước batch dùng cho warmup sẽ chạy thử một lần và capture CUDA graphs. CUDA graph ghi lại toàn bộ chuỗi công việc của GPU thành một DAG. Sau này trong forward pass, chúng ta chỉ cần launch/replay các graph đã được chuẩn bị sẵn, nhờ đó giảm overhead khởi chạy kernel và cải thiện độ trễ.

Bây giờ khi engine đã được khởi tạo, hãy chuyển sang hàm generate.

2. Hàm Generate

Bước đầu tiên là xác thực và đưa các request vào engine. Với mỗi prompt, chúng ta:

Tạo một request ID duy nhất và ghi nhận thời điểm request đến.
Gọi một bộ tiền xử lý đầu vào để tokenize prompt và trả về một dictionary chứa prompt, prompt_token_ids và type (text, tokens, embeds, v.v.).
Đóng gói thông tin này vào một EngineCoreRequest, đồng thời thêm priority, sampling params và các metadata khác.
Chuyển request vào engine core; engine sẽ bọc nó trong một đối tượng Request và đặt trạng thái thành WAITING. Request này sau đó được thêm vào hàng đợi waiting của scheduler (append nếu dùng FCFS, hoặc heap-push nếu dùng priority).

Đến thời điểm này, engine đã được nạp request và việc thực thi có thể bắt đầu. Trong ví dụ engine đồng bộ, các prompt ban đầu này là những request duy nhất mà chúng ta sẽ xử lý – không có cơ chế đưa thêm request mới vào giữa lúc đang chạy. Ngược lại, engine bất đồng bộ hỗ trợ điều này (còn gọi là continuous batching [6]): sau mỗi bước, cả request mới lẫn request cũ đều được xem xét.

Lưu ý: Vì forward pass làm phẳng toàn bộ batch thành một chuỗi duy nhất và các custom kernel xử lý nó một cách hiệu quả, nên continuous batching về bản chất đã được hỗ trợ ngay cả trong engine đồng bộ.

Tiếp theo, miễn là vẫn còn request cần xử lý, engine sẽ liên tục gọi hàm step(). Mỗi bước gồm ba giai đoạn:

Schedule: chọn những request sẽ chạy ở bước này (decode và/hoặc prefill (chunked)).
Forward pass: chạy mô hình và sample token.
Postprocess: thêm các token ID đã sample vào từng Request, detokenize và kiểm tra điều kiện dừng. Nếu một request đã hoàn thành, thực hiện dọn dẹp (ví dụ trả các block KV cache của nó về free_block_queue) và trả kết quả sớm.

Lưu ý: Các điều kiện dừng là:

Request vượt quá giới hạn độ dài của nó (max_model_length hoặc max_tokens riêng của request).
Token được sample là EOS ID (trừ khi bật ignore_eos — hữu ích cho benchmark khi muốn ép mô hình sinh ra một số lượng token đầu ra nhất định).
Token được sample khớp với bất kỳ stop_token_ids nào được chỉ định trong sampling parameters.
Xuất hiện stop string trong đầu ra — chúng ta sẽ cắt đầu ra đến vị trí xuất hiện stop string đầu tiên và hủy request trong engine (lưu ý rằng stop_token_ids sẽ xuất hiện trong đầu ra, còn stop string thì không).

Vòng lặp của engine

Lưu ý: Trong chế độ streaming, chúng ta sẽ gửi các token trung gian ngay khi chúng được sinh ra, nhưng hiện tại hãy bỏ qua điều đó.

Tiếp theo, chúng ta sẽ xem xét việc lập lịch (scheduling) chi tiết hơn.

Scheduler

Có hai loại workload chính mà một inference engine phải xử lý:

Prefill request — một forward pass trên toàn bộ các token của prompt. Chúng thường bị giới hạn bởi khả năng tính toán (compute-bound) (ngưỡng cụ thể phụ thuộc vào phần cứng và độ dài prompt). Ở cuối bước này, chúng ta sample một token duy nhất từ phân phối xác suất tại vị trí token cuối cùng.
Decode request — một forward pass chỉ trên token mới nhất. Tất cả các vector KV trước đó đã được cache sẵn. Chúng bị giới hạn bởi băng thông bộ nhớ (memory-bandwidth-bound), vì chúng ta vẫn phải nạp toàn bộ trọng số của LLM (và KV cache) chỉ để tính ra một token.

Lưu ý: Trong phần benchmark chúng ta sẽ phân tích mô hình “roofline” của hiệu năng GPU. Phần đó sẽ đi sâu hơn vào đặc điểm hiệu năng của prefill và decode.

Scheduler của V1 có thể trộn cả hai loại request trong cùng một bước, nhờ các lựa chọn thiết kế thông minh hơn. Ngược lại, engine V0 chỉ có thể xử lý hoặc prefill hoặc decode tại một thời điểm.

Scheduler ưu tiên các request decode — tức là những request đã nằm trong hàng đợi running.

Với mỗi request như vậy, nó sẽ:

Tính số lượng token mới cần sinh ra (không phải lúc nào cũng là 1, do speculative decoding và async scheduling — sẽ nói sau).
Gọi hàm allocate_slots của KV-cache manager (chi tiết bên dưới).
Cập nhật token budget bằng cách trừ đi số token ở bước 1.

Sau đó, nó xử lý các request prefill trong hàng đợi waiting:

Lấy số lượng block đã được tính toán (trả về 0 nếu prefix caching bị tắt — sẽ nói sau).
Gọi hàm allocate_slots của KV-cache manager.
Lấy request ra khỏi waiting và chuyển sang running, đồng thời đặt trạng thái thành RUNNING.
Cập nhật token budget.

Bây giờ hãy xem allocate_slots thực hiện những gì:

Tính số lượng block – xác định cần cấp phát bao nhiêu block KV cache mới (n). Mỗi block mặc định lưu 16 token. Ví dụ, nếu một request prefill có 17 token mới, chúng ta cần ceil(17/16) = 2 block.
Kiểm tra khả dụng – nếu không còn đủ block trong pool của manager, hàm sẽ thoát sớm. Tùy request là decode hay prefill, engine có thể thử recompute preemption (V0 còn hỗ trợ swap preemption) bằng cách đẩy các request ưu tiên thấp ra ngoài (gọi kv_cache_manager.free để trả các block KV về pool), hoặc bỏ qua việc schedule request đó và tiếp tục thực thi.
Cấp phát block – thông qua coordinator của KV-cache manager, lấy n block đầu tiên từ block pool (danh sách liên kết kép free_block_queue đã nhắc trước đó). Sau đó lưu vào req_to_blocks, dictionary ánh xạ mỗi request_id tới danh sách các block KV cache của nó.

Danh sách các block KV cache

Cuối cùng, chúng ta đã sẵn sàng để thực hiện một forward pass!

Chạy forward pass

Chúng ta gọi execute_model của model executor, hàm này ủy quyền cho Worker, và Worker tiếp tục ủy quyền cho model runner.

Đây là các bước chính:

Cập nhật trạng thái — loại bỏ các request đã hoàn thành khỏi input_batch; cập nhật các metadata liên quan đến forward pass (ví dụ: các block KV cache của từng request sẽ được dùng để đánh chỉ mục vào bộ nhớ paged KV cache).
Chuẩn bị đầu vào — sao chép các buffer từ CPU → GPU; tính toán vị trí (positions); xây dựng slot_mapping (sẽ nói rõ hơn trong ví dụ); tạo metadata cho attention.
Forward pass — chạy mô hình bằng các kernel paged attention tùy biến. Tất cả các sequence được làm phẳng và nối lại thành một “super sequence” dài. Chỉ số vị trí và attention mask đảm bảo mỗi sequence chỉ attention tới các token của chính nó, nhờ đó cho phép continuous batching mà không cần right-padding.
Thu thập trạng thái của token cuối — trích xuất hidden state tại vị trí cuối của mỗi sequence và tính toán logits.
Sample — sample token từ các logits đã tính, theo cấu hình sampling (greedy, temperature, top-p, top-k, v.v.).

Bản thân bước forward pass có hai chế độ thực thi:

Eager mode — chạy forward pass chuẩn của PyTorch khi bật eager execution.
“Captured” mode — thực thi / phát lại một CUDA Graph đã được capture trước khi không ép dùng eager (hãy nhớ rằng chúng ta đã capture các graph này trong lúc khởi tạo engine ở bước khởi tạo KV cache).

Đây là một ví dụ cụ thể giúp làm rõ continuous batching và paged attention: