Chunked Prefill

Prefill theo khối là kỹ thuật phục vụ model, chia prompt đầu vào dài thành các nhóm token nhỏ hơn và xử lý qua nhiều bước, xen kẽ với việc tạo token cho các yêu cầu khác. Nếu không có kỹ thuật này, một engine suy luận tiếp nhận prompt dài sẽ dành trọn một bước cho prompt đó, khiến mọi phản hồi đang được truyền đến người dùng khác tạm dừng cho đến khi xử lý xong. Bài báo SARATHI năm 2023 giới thiệu prefill theo khối, còn nghiên cứu tiếp theo Sarathi-Serve xây dựng phương pháp lập lịch “không gián đoạn” dựa trên kỹ thuật này. Dynamic SplitFuse của DeepSpeed-FastGen dùng cùng ý tưởng, và hiện prefill theo khối được bật mặc định trong vLLM, đồng thời có trong SGLang và TensorRT-LLM, những engine phục vụ nhiều mô hình ngôn ngữ lớn (LLM) có trọng số mở.

Prefill, decode và tình trạng gián đoạn sinh token#

Suy luận Transformer tự hồi quy có hai giai đoạn:

• Prefill: Model xử lý song song mọi token trong prompt, tính key và value của chúng rồi lưu vào KV cache. Prefill thực hiện nhiều phép toán trên mỗi byte trọng số được đọc, nên thường bị giới hạn bởi năng lực tính toán.
• Decode: Model tạo một token mỗi bước cho từng yêu cầu, mỗi lần đều đọc toàn bộ cache. Decode thực hiện ít phép toán trên mỗi byte được đọc, nên thường bị giới hạn bởi băng thông bộ nhớ.

Một bước GPU chứa prefill 20.000 token mất nhiều thời gian hơn hẳn một bước decode, khiến mọi yêu cầu trong batch phải chờ. Sarathi-Serve gọi những khoảng dừng này là tình trạng gián đoạn sinh token; với người dùng có prompt ngắn, chúng xuất hiện dưới dạng các đỉnh độ trễ giữa các token.

Cách thức hoạt động của chunked prefill#

Bộ lập lịch cấp cho mỗi bước một ngân sách token, được đặt trong vLLM bằng max_num_batched_tokens. Trước tiên, bộ lập lịch thêm một token cho mỗi lượt decode đang chạy, sau đó dùng phần ngân sách còn lại cho các khối prompt đang chờ. Nếu không vừa, prompt sẽ tiếp tục ở bước kế tiếp, mở rộng KV cache theo từng khối cho đến khi xử lý xong toàn bộ prompt và bắt đầu decode. Kết hợp hai giai đoạn trong một bước cũng giúp sử dụng GPU hiệu quả hơn, vì decode để phần năng lực tính toán nhàn rỗi mà các khối prefill có thể tận dụng:

Chia khối chỉ thay đổi thời điểm xử lý token của prompt. Mỗi khối chú ý đến key và value đã lưu trong cache của mọi khối trước đó cùng các token của chính nó theo causal mask, nên kết quả giống với prefill một lượt, như minh họa trong một lớp attention của PyTorch:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
prompt_len, chunk = 1000, 256
q, k, v = torch.randn(3, 1, 8, prompt_len, 64).unbind(0)  # one layer: 8 heads, 64 dims each

# One-shot prefill: all prompt tokens attend causally in a single pass
full = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# Chunked prefill: each chunk attends to the cached keys and values of earlier chunks plus itself
outputs = []
for start in range(0, prompt_len, chunk):
    end = min(start + chunk, prompt_len)
    mask = torch.ones(end - start, end, dtype=torch.bool).tril(diagonal=start)  # causal over absolute positions
    q_chunk, k_cached, v_cached = q[..., start:end, :], k[..., :end, :], v[..., :end, :]
    outputs.append(F.scaled_dot_product_attention(q_chunk, k_cached, v_cached, attn_mask=mask))
chunked = torch.cat(outputs, dim=2)

print(f"{len(outputs)} chunks, same result as one pass: {torch.allclose(full, chunked, atol=1e-5)}")

Lệnh này in ra 4 chunks, same result as one pass: True: prompt 1.000 token được xử lý thành các khối gồm 256, 256, 256 và 232 token, và đầu ra khớp với một lượt xử lý.

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

• Batch liên tục: Thêm và loại bỏ yêu cầu ở mỗi bước thay vì chờ cả batch hoàn tất. Prefill theo khối giúp lập lịch cho prompt dài bên trong vòng lặp đó.
• PagedAttention: Lưu KV cache trong các trang có kích thước cố định, nhờ đó prompt được prefill qua nhiều bước có thể mở rộng cache mà không cần cấp phát trước bộ nhớ cho toàn bộ độ dài prompt.
• Bộ nhớ đệm prompt: Tái sử dụng key và value đã lưu trong cache của tiền tố prompt lặp lại, qua đó loại bỏ hoàn toàn phần prefill tương ứng.
• Tách biệt prefill và decode: Chạy hai giai đoạn trên các cụm GPU riêng và truyền KV cache giữa chúng, để mỗi cụm có thể được mở rộng và tinh chỉnh theo điểm nghẽn riêng.
• Decode suy đoán: Tăng tốc giai đoạn decode bằng cách soạn trước một số token rồi xác minh chúng trong một lượt.

Ứng dụng#

Prefill theo khối đặc biệt hữu ích khi prompt dài và ngắn cùng dùng chung một server. Trợ lý tạo sinh tăng cường truy xuất (RAG) đưa báo cáo 50 trang vào cửa sổ ngữ cảnh có thể prefill theo từng khối trong khi các cuộc trò chuyện của người dùng khác vẫn tiếp tục được truyền. Model thị giác-ngôn ngữ (VLM) bổ sung hàng trăm hoặc hàng nghìn token ảnh vào mỗi prompt, vì vậy pipeline camera gửi các khung hình để mô tả sẽ tạo ra chính xác kiểu tải hỗn hợp đó. Giao diện LLM của Ultralytics gửi văn bản và ảnh đến mọi endpoint tương thích OpenAI, bao gồm server vLLM tự lưu trữ; tại đó, các yêu cầu từ pipeline YOLO26 mô tả những khung hình phát hiện được người sẽ được prefill theo khối song song với lưu lượng của người dùng khác.

Tinh chỉnh ngân sách token#

Ngân sách tạo ra sự đánh đổi giữa hai chỉ số độ trễ. Ngân sách nhỏ dành ít chỗ cho prefill ở mỗi bước, giúp độ trễ giữa các token thấp đối với người dùng nhận kết quả liên tục nhưng làm chậm thời gian đến token đầu tiên của prompt mới; ngân sách lớn thì ngược lại. Mỗi khối tiếp theo của prompt dài cũng phải đọc lại KV cache đang tăng dần của các khối trước, nên khối quá nhỏ làm tăng lưu lượng bộ nhớ. Người vận hành tinh chỉnh ngân sách dựa trên độ trễ suy luận trung vị và ở phần đuôi, thời gian đến token đầu tiên và thông lượng trên độ dài prompt thực tế. Hướng dẫn tối ưu hóa vLLM nêu 2.048 token là mức ngân sách ưu tiên độ trễ giữa các token và khuyến nghị hơn 8.192 token để đạt thông lượng tối đa với model nhỏ hơn trên GPU lớn.