Bài trước, chúng ta đã tìm hiểu về LLM Engine & Engine Core trong hệ thống suy luận LLM thông lượng cao. Bài này, chúng ta sẽ đi sâu vào:

• Chunked prefill
• Prefix caching
• Guided decoding (thông qua các máy trạng thái hữu hạn bị ràng buộc bởi ngữ pháp)
• Speculative decoding
• Disaggregated P/D (prefill/decoding)

Chunked prefill là một kỹ thuật để xử lý các prompt dài bằng cách chia bước prefill của chúng thành các phần nhỏ hơn. Nếu không có nó, chúng ta có thể gặp tình huống một yêu cầu rất dài chiếm toàn bộ một bước của engine, khiến các yêu cầu prefill khác không thể chạy. Điều đó sẽ trì hoãn tất cả các yêu cầu khác và làm tăng độ trễ của chúng.

Ví dụ, giả sử mỗi chunk chứa n (=8) token, được gắn nhãn bằng các chữ cái thường và ngăn cách bằng dấu “-“. Một prompt dài P có thể trông như x-y-z, trong đó z là một chunk chưa hoàn chỉnh (ví dụ: 2 token). Khi đó, thực thi toàn bộ prefill cho P sẽ mất ≥ 3 bước của engine (> có thể xảy ra nếu nó không được lập lịch để thực thi ở một trong các bước), và chỉ ở bước chunked prefill cuối cùng chúng ta mới lấy mẫu một token mới.

Đây là cùng ví dụ đó dưới dạng trực quan:

Việc triển khai khá đơn giản: đặt giới hạn số lượng token mới trong mỗi bước. Nếu số lượng được yêu cầu vượt quá long_prefill_token_threshold, hãy đặt lại nó chính xác bằng giá trị đó. Logic lập chỉ mục nền tảng (đã mô tả trước đó) sẽ xử lý phần còn lại.

Trong vLLM V1, bạn bật chunked prefill bằng cách đặt long_prefill_token_threshold thành một số nguyên dương. (Về mặt kỹ thuật, điều này vẫn có thể xảy ra ngay cả khi không đặt giá trị này; nếu độ dài prompt vượt quá ngân sách token, chúng ta sẽ cắt bớt nó và chạy chunked prefill.)

Để giải thích cách hoạt động của prefix caching, hãy lấy ví dụ mã gốc và chỉnh sửa nó một chút:

Prefix caching tránh việc tính toán lại các token mà nhiều prompt cùng chia sẻ ở phần đầu – vì thế mới gọi là prefix (tiền tố).

Thành phần quan trọng nhất là long_prefix: nó được định nghĩa là bất kỳ tiền tố nào dài hơn một block của KV cache (mặc định là 16 token). Để đơn giản hóa ví dụ, hãy giả sử long_prefix có độ dài chính xác bằng n x block_size (trong đó n ≥ 1).

Lưu ý: tức là nó khớp hoàn hảo với ranh giới block, nếu không, chúng ta sẽ phải tính toán lại `long_prefix_len % block_size` token vì không thể lưu cache cho các block chưa hoàn chỉnh.

Nếu không có prefix caching, mỗi lần xử lý một request mới có cùng long_prefix, chúng ta sẽ phải tính toán lại toàn bộ n x block_size token.

Khi có prefix caching, các token đó chỉ được tính toán một lần (các KV của chúng được lưu trong bộ nhớ phân trang của KV cache) rồi được tái sử dụng, vì vậy chỉ những token mới của prompt cần được xử lý. Điều này giúp tăng tốc các request ở giai đoạn prefill (mặc dù không giúp ích cho giai đoạn decode).

Nó hoạt động như thế nào trong vLLM?

Trong lần gọi generate đầu tiên, ở giai đoạn lập lịch (scheduling), bên trong kv_cache_manager.get_computed_blocks, engine sẽ gọi hash_request_tokens:

• Hàm này chia long_prefix + prompts[0] thành các đoạn 16 token.
• Với mỗi đoạn hoàn chỉnh, nó tính một giá trị hash (sử dụng hash tích hợp sẵn hoặc SHA-256, chậm hơn nhưng ít va chạm hơn). Giá trị hash được kết hợp từ hash của block trước đó, các token hiện tại và metadata tùy chọn.

Lưu ý: metadata tùy chọn bao gồm: MM hash, LoRA ID, cache salt (được chèn vào hash của block đầu tiên để đảm bảo chỉ những request có cùng cache salt mới có thể tái sử dụng các block).

• Mỗi kết quả được lưu dưới dạng một đối tượng BlockHash, chứa cả giá trị hash và các token ID của nó. Chúng ta trả về một danh sách các block hash.

Danh sách này được lưu trong self.req_to_block_hashes[request_id].

Tiếp theo, engine gọi find_longest_cache_hit để kiểm tra xem có hash nào trong số này đã tồn tại trong cached_block_hash_to_block hay chưa. Ở request đầu tiên, sẽ không tìm thấy block nào trùng khớp.

Sau đó chúng ta gọi allocate_slots, hàm này sẽ gọi coordinator.cache_blocks, hàm này liên kết các mục BlockHash mới với các KV block đã được cấp phát và ghi nhận chúng vào cached_block_hash_to_block.

Sau đó, forward pass sẽ điền các KV vào vùng nhớ KV cache phân trang tương ứng với các KV cache block mà chúng ta đã cấp phát ở trên.

Lưu ý: Sau nhiều bước của engine, nó sẽ cấp phát thêm các KV cache block, nhưng điều đó không quan trọng đối với ví dụ này vì tiền tố đã phân nhánh ngay sau `long_prefix`.

Ở lần gọi generate thứ hai với cùng tiền tố, các bước 1–3 được lặp lại, nhưng lúc này find_longest_cache_hit sẽ tìm thấy các block khớp với toàn bộ n block (thông qua tìm kiếm tuyến tính). Engine có thể tái sử dụng trực tiếp các KV block đó.

Nếu request gốc vẫn còn tồn tại, bộ đếm tham chiếu (reference count) của các block đó sẽ tăng lên (ví dụ thành 2). Trong ví dụ này, request đầu tiên đã hoàn thành, nên các block đã được giải phóng trở lại pool và reference count của chúng được đặt lại về 0. Vì chúng ta có thể lấy lại chúng từ cached_block_hash_to_block, nên chúng ta biết rằng chúng vẫn hợp lệ (logic của KV cache manager được thiết kế theo cách đó), vì vậy chúng ta chỉ cần xóa chúng khỏi free_block_queue một lần nữa.

Lưu ý nâng cao:KV-cache block chỉ trở nên không hợp lệ khi chúng sắp được cấp phát lại từ free_block_queue (hàng đợi lấy phần tử từ bên trái) và chúng ta phát hiện block đó vẫn còn có hash liên kết và vẫn hiện diện trong cached_block_hash_to_block. Tại thời điểm đó, chúng ta sẽ xóa hash của block và loại bỏ mục tương ứng khỏi cached_block_hash_to_block, đảm bảo rằng block đó không thể được tái sử dụng thông qua prefix caching nữa (ít nhất là không cho tiền tố cũ đó).

Và đó chính là cốt lõi của prefix caching: đừng tính toán lại những tiền tố mà bạn đã thấy trước đó — chỉ cần tái sử dụng KV cache của chúng!

Nếu bạn hiểu ví dụ này thì bạn cũng đã hiểu cách hoạt động của paged attention.

Prefix caching được bật mặc định. Để tắt nó: enable_prefix_caching = False.

Guided decoding là một kỹ thuật mà ở mỗi bước giải mã (decoding), các logits sẽ bị ràng buộc bởi một máy trạng thái hữu hạn (finite state machine) dựa trên ngữ pháp (grammar). Điều này đảm bảo rằng chỉ những token được ngữ pháp cho phép mới có thể được lấy mẫu.

Đây là một cơ chế rất mạnh: bạn có thể áp đặt mọi thứ, từ ngữ pháp chính quy (regular grammar, Chomsky loại 3, ví dụ các mẫu regex tùy ý) cho đến ngữ pháp phi ngữ cảnh (context-free grammar, loại 2, bao phủ hầu hết các ngôn ngữ lập trình).

Để làm điều này bớt trừu tượng, hãy bắt đầu với ví dụ đơn giản nhất có thể, dựa trên đoạn mã trước đó của chúng ta:

Trong ví dụ đồ chơi đã đưa ra (giả sử tokenization ở mức ký tự): ở giai đoạn prefill, FSM sẽ che (mask) các logits để chỉ còn “P” hoặc “N” là có thể được chọn. Nếu “P” được lấy mẫu, FSM sẽ chuyển sang nhánh “Positive”; ở bước tiếp theo chỉ “o” được phép, và cứ tiếp tục như vậy.

Cách điều này hoạt động trong vLLM:

1. Khi khởi tạo LLM engine, một StructuredOutputManager được tạo ra; nó có quyền truy cập vào tokenizer và duy trì một tensor _grammar_bitmask.
2. Khi thêm một request, trạng thái của nó được đặt thành WAITING_FOR_FSM và grammar_init sẽ chọn backend compiler (ví dụ: xgrammar [7]; lưu ý rằng các backend là mã của bên thứ ba).
3. Ngữ pháp cho request này được biên dịch bất đồng bộ.
4. Trong quá trình scheduling, nếu việc biên dịch bất đồng bộ đã hoàn tất, trạng thái sẽ chuyển sang WAITING và request_id được thêm vào structured_output_request_ids; nếu chưa, nó sẽ được đưa vào skipped_waiting_requests để thử lại ở bước engine tiếp theo.
5. Sau vòng lặp scheduling (vẫn còn bên trong scheduling), nếu có các request FSM, StructuredOutputManager sẽ yêu cầu backend chuẩn bị/cập nhật _grammar_bitmask.
6. Sau khi forward pass tạo ra logits, hàm của xgr_torch_compile sẽ mở rộng bitmask thành kích thước của vocabulary (tỷ lệ mở rộng 32× vì chúng ta dùng số nguyên 32 bit) và che các logits không được phép bằng –∞.
7. Sau khi lấy mẫu token tiếp theo, FSM của request sẽ được tiến lên thông qua accept_tokens. Về mặt trực quan, chúng ta chuyển sang trạng thái tiếp theo trên sơ đồ FSM.

Trong quá trình sinh tự hồi quy (autoregressive generation), mỗi token mới đều cần một lần forward pass của mô hình ngôn ngữ lớn (large LM). Điều này rất tốn kém – ở mỗi bước, toàn bộ trọng số của mô hình đều phải được nạp lại và áp dụng chỉ để tính ra một token duy nhất! (giả sử batch size = 1, tổng quát là B).

Speculative decoding tăng tốc quá trình này bằng cách đưa vào một mô hình nháp (draft LM) nhỏ hơn. Mô hình nháp sẽ đề xuất k token với chi phí thấp. Nhưng mục tiêu cuối cùng không phải là lấy mẫu từ mô hình nhỏ — nó chỉ dùng để đoán các phần tiếp theo có thể xảy ra. Mô hình lớn vẫn là bên quyết định điều gì là hợp lệ.

Các bước như sau:

• Draft: chạy mô hình nhỏ trên ngữ cảnh hiện tại và đề xuất k token.

• Verify: chạy mô hình lớn một lần trên ngữ cảnh + k token nháp. Điều này tạo ra xác suất cho k vị trí đó cộng thêm một vị trí nữa (vì vậy ta có k+1 ứng viên).

• Accept/reject: duyệt từ trái sang phải qua k token nháp:

• • • Nếu xác suất của mô hình lớn cho token nháp ≥ xác suất của mô hình nháp, chấp nhận token đó.

• Nếu không, chấp nhận nó với xác suất p_large(token) / p_draft(token).

• Dừng ở lần từ chối đầu tiên, hoặc chấp nhận toàn bộ k token nháp.

• Nếu toàn bộ k token nháp được chấp nhận, ta cũng lấy mẫu token thứ (k+1) “miễn phí” từ mô hình lớn (vì phân phối đó đã được tính rồi).

• Nếu có một token bị từ chối, tạo một phân phối cân bằng lại mới tại vị trí đó (p_large – p_draft, chặn giá trị nhỏ nhất ở 0, chuẩn hóa để tổng bằng 1) và lấy mẫu token cuối cùng từ phân phối này.

Tại sao cách này hoạt động: Mặc dù chúng ta dùng mô hình nhỏ để đề xuất ứng viên, quy tắc chấp nhận/từ chối đảm bảo rằng về kỳ vọng, chuỗi kết quả có phân phối chính xác giống như khi lấy mẫu từng token một từ mô hình lớn. Điều này có nghĩa là speculative decoding tương đương về mặt thống kê với decoding tự hồi quy tiêu chuẩn — nhưng có thể nhanh hơn đáng kể, vì một lần chạy mô hình lớn có thể tạo ra tối đa k+1 token.

vLLM V1 không hỗ trợ phương pháp dùng một LLM nhỏ làm draft model; thay vào đó nó triển khai các cơ chế đề xuất nhanh hơn nhưng kém chính xác hơn: n-gram, EAGLE [9], và Medusa [10].

Tóm tắt ngắn gọn từng phương pháp:

• n-gram: lấy prompt_lookup_max token cuối cùng; tìm một lần xuất hiện trước đó trong chuỗi; nếu tìm thấy, đề xuất k token đã xuất hiện sau lần khớp đó; nếu không, giảm kích thước cửa sổ và thử lại cho đến prompt_lookup_min.

• • EAGLE: thực hiện “phẫu thuật mô hình” trên mô hình lớn – giữ nguyên embedding và LM head, thay toàn bộ transformer stack bằng một MLP nhẹ; sau đó fine-tune phần này như một draft model giá rẻ.

• Medusa: huấn luyện các đầu tuyến tính phụ trợ đặt trên mô hình lớn (ở các embedding trước LM head) để dự đoán song song k token tiếp theo; dùng các head này để đề xuất token hiệu quả hơn so với chạy một mô hình nhỏ riêng biệt.

Dưới đây là cách gọi speculative decoding trong vLLM bằng cách dùng ngram làm phương pháp draft:

Nó hoạt động như thế nào trong vLLM?

Thiết lập (trong quá trình khởi tạo engine)

1. Khởi tạo thiết bị: tạo một drafter (mô hình nháp, ví dụ NgramProposer) và một rejection_sampler (một phần của nó được viết bằng Triton).
2. Tải mô hình: tải trọng số của draft model (không làm gì đối với n-gram).

Sau đó trong hàm generate (giả sử nhận một request hoàn toàn mới)

1. Chạy bước prefill thông thường bằng mô hình lớn.
2. Sau forward pass và lấy mẫu tiêu chuẩn, gọi propose_draft_token_ids(k) để lấy mẫu k token nháp từ draft model.
3. Lưu các token này vào request.spec_token_ids (cập nhật metadata của request).
4. Ở bước engine tiếp theo, khi request nằm trong hàng đợi đang chạy, cộng len(request.spec_token_ids) vào số lượng “token mới” để allocate_slots đặt trước đủ KV block cho lần forward pass.
5. Sao chép spec_token_ids vào input_batch.token_ids_cpu để tạo thành các token (context + draft).
6. Tính metadata thông qua _calc_spec_decode_metadata (hàm này sao chép các token từ input_batch.token_ids_cpu, chuẩn bị logits, v.v.), sau đó chạy một forward pass của mô hình lớn trên các token nháp.
7. 1. Thay vì lấy mẫu thông thường từ logits, sử dụng rejection_sampler để chấp nhận/từ chối theo thứ tự từ trái sang phải và tạo ra output_token_ids.
   2. Lặp lại các bước 2–7 cho đến khi đạt điều kiện dừng.

   Cách tốt nhất để thực sự hiểu điều này là mở debugger và bước từng dòng qua codebase, nhưng phần này hy vọng đã cho bạn một cảm nhận tổng quát về cơ chế hoạt động của nó.

Prefill và decode có đặc tính hiệu năng rất khác nhau (compute-bound so với memory-bandwidth-bound), vì vậy việc tách riêng quá trình thực thi của chúng là một thiết kế hợp lý. Điều này cho phép kiểm soát độ trễ chặt chẽ hơn – cả TFTT (time-to-first-token) và ITL (inter-token latency).

Trong thực tế, chúng ta chạy N instance vLLM cho prefill và M instance vLLM cho decode, rồi tự động co giãn (autoscaling) chúng dựa trên hỗn hợp request đang có tại thời điểm thực. Các worker prefill sẽ ghi KV vào một dịch vụ KV-cache chuyên dụng; các worker decode sẽ đọc từ đó. Điều này giúp cô lập phần prefill dài và bùng nổ khỏi phần decode ổn định và nhạy cảm với độ trễ.

Điều này hoạt động như thế nào trong vLLM?

Để rõ ràng, ví dụ dưới đây sử dụng SharedStorageConnector, một triển khai connector dùng cho mục đích gỡ lỗi nhằm minh họa cơ chế hoạt động.

Lưu ý:Connector là lớp trừu tượng của vLLM để xử lý việc trao đổi KV giữa các instance. Giao diện connector hiện chưa ổn định, và có một số cải tiến ngắn hạn đang được lên kế hoạch, có thể bao gồm cả những thay đổi phá vỡ tương thích.

Chúng ta khởi chạy 2 instance vLLM (GPU 0 cho prefill và GPU 1 cho decode), sau đó chuyển KV cache giữa chúng:

Đây là các bước trong vLLM:

1. Khởi tạo (Instantiation) – Trong quá trình xây dựng engine, các connector được tạo ở hai nơi:
   • Bên trong thủ tục khởi tạo thiết bị của worker (trong hàm khởi tạo môi trường phân tán của worker), với vai trò “worker”.
   • Bên trong constructor của scheduler, với vai trò “scheduler”.
2. Tra cứu cache (Cache lookup) – Khi scheduler xử lý các request prefill từ hàng đợi waiting (sau khi kiểm tra prefix-cache cục bộ), nó gọi connector.get_num_new_matched_tokens. Hàm này kiểm tra các token đã được cache bên ngoài trong máy chủ KV-cache. Prefill luôn nhận 0 ở đây; decode có thể có cache hit. Kết quả này được cộng vào số lượng cục bộ trước khi gọi allocate_slots.
3. Cập nhật trạng thái (State update) – Sau đó scheduler gọi connector.update_state_after_alloc, hàm này ghi nhận các request có cache (không làm gì đối với prefill).
4. Xây dựng đối tượng metadata (Build metadata object) – Ở cuối giai đoạn scheduling, scheduler gọi meta = connector.build_connector_meta:
   • Prefill thêm tất cả các request với is_store=True (để tải KV lên).
   • Decode thêm các request với is_store=False (để lấy KV xuống).
5. Context manager — Trước khi forward pass diễn ra, engine đi vào một context manager của KV-connector:
   • Khi vào: gọi kv_connector.start_load_kv. Đối với decode, hàm này tải KV từ máy chủ bên ngoài và chèn nó vào bộ nhớ phân trang. Đối với prefill, không làm gì.
   • Khi thoát: gọi kv_connector.wait_for_save. Đối với prefill, hàm này sẽ chặn cho đến khi KV được tải lên máy chủ bên ngoài. Đối với decode, không làm gì.

Dưới đây là một ví dụ trực quan:

Lưu ý bổ sung:

• Đối với SharedStorageConnector, “máy chủ bên ngoài” thực chất chỉ là hệ thống tệp cục bộ (local file system).

• Tùy theo cấu hình, việc truyền KV cũng có thể được thực hiện theo từng lớp (layer-by-layer), tức là trước hoặc sau mỗi lớp attention.

• Decode chỉ tải KV từ bên ngoài một lần duy nhất, ở bước đầu tiên của request; sau đó nó sẽ tự tính toán và lưu trữ cục bộ.

Phần 1: LLM Engine & Engine Core

Phần 3: Scaling up: từ một GPU đến nhiều GPU