Bước inference: xử lý prompt và chọn token đầu tiên

1Định nghĩa

Prefill là pha đầu tiên của suy luận sau khi bộ lập lịch cấp lượt chạy cho mảng token từ bước tách token: worker đưa toàn bộ N token của prompt qua các lớp Transformer trong cùng một lượt tính song song. Vì mọi token đầu vào đã biết trước, không cần chờ nhau, và chính các phép tính đó là lượt forward quen thuộc ở cột 2 của sơ đồ, chỉ khác là trọng số đã đóng băng.

Tại mỗi lớp attention, cặp Key và Value của cả N vị trí được ghi vào KV cache trên HBM. Riêng bảng logits ở vị trí cuối prompt được chuyển sang bộ chọn token (sampling) để chọn token đầu tiên của câu trả lời, khép lại thời gian chờ token đầu (TTFT) và mở màn vòng decode.

Ví dụ đời thường: Prefill giống khoảnh khắc người phiên dịch nghe trọn vẹn câu hỏi của khách: nắm bắt toàn bộ ngữ cảnh trong một lần, ghi nhanh các ý chính ra sổ tay (KV cache) rồi bật ra ngay chữ đầu tiên của câu trả lời. Khác người phiên dịch, mô hình xử lý mọi từ của câu hỏi cùng lúc bằng phép nhân ma trận chứ không đọc tuần tự từng từ.

2Luồng

Đây là nửa đầu của ô 5B, ‘Prefill & decode’. Mũi tên vào mang N token ID; mũi tên ra mang hai thứ: KV cache (đi vào vòng decode) và một token đã chọn. Phép tính bên trong chính là cột 2 của sơ đồ, nhưng không có loss, backward hay cập nhật.

Sơ đồ kiến trúc huấn luyện · 22 ô (1A → 5C)

Kiến trúc huấn luyện LLM · Ánh xạ với Bài Lab tổng hợp

Thứ tự đọc: 1A → 1B ở cột 1, 2A → 2D ở cột 2, 3A → 3E ở cột 3; mũi tên next batch từ 3E quay về cột 1 và mũi tên load weights đi xuống băng 5 (5A → 5C). Cột 4 không phải bước kế tiếp: 4A, 4B tiết kiệm bộ nhớ, còn 4C–4G là các phương pháp huấn luyện thay thế nhau. Di chuột hoặc bấm vào từng ô để đọc pop-over.

Tất cả (22)1. Dữ liệu & nhãn2. Forward3. Học từ sai số4. Quy mô & phương pháp5. Suy luận

Ẩn khung overlay

1A1B1C2A2B2C2D3A3B3C3D3E4A4B4C4D4E4F4G5A5B5C

5B5 · Checkpoint → Inference

prompt một lượt · rồi mỗi bước một token · KV cache

5B · Prefill & Decode

Prefill đọc cả prompt một lượt, rồi decode sinh từng token với KV cache.

Nhận gì

Dãy token ID của prompt từ ô 5A, cùng trọng số đã nạp từ checkpoint ở ô 3E.

Làm gì, vì sao

Prefill đưa cả prompt qua mô hình một lượt (lặp lại cột 2) để lấp đầy KV cache (bộ nhớ key và value của mọi vị trí) và lấy logits ở vị trí cuối. Decode mỗi bước chỉ đưa một token mới, đọc lại K, V cũ từ cache, thêm K, V mới rồi chọn token kế; nhờ cache không phải tính lại cả lịch sử.

Ra gì

Từng token được chọn, vừa nối vào chuỗi vừa đi sang ô 5C; vòng lặp tiếp tục đến khi gặp điều kiện dừng.

Ví dụ, cách đọc ô

Nhãn ô: prompt once · then one token per step · KV cache. Ví dụ trong lab: prompt 13 token và đầu ra 5 token cần 13 + 4 = 17 vị trí xử lý với cache, so với 75 nếu mỗi bước tính lại (số vị trí, không phải thời gian).

Lưu ý

Cache không đổi kết quả toán học, chỉ tránh tính lại. Thời gian thật phụ thuộc phần cứng, và cache lớn dần theo độ dài chuỗi.

Đọc bài 5B · Prefill & Decode↗Bài liên quan · Bước inference: vòng sinh và bổ sung context↗Giai đoạn 13 trong Bài Lab: Từ checkpoint sang suy luận↗

flowchart TD
  step1["Prompt đã thành token ID"]
  step2["Forward qua các lớp"]
  step3["Ghi Key và Value"]
  step4["Chọn token đầu"]
  step1 --> step2
  step2 --> step3
  step3 --> step4

3Giải pháp

Vị trí trên sơ đồ: nửa đầu của ô 5B "Prefill & decode" trong băng "Checkpoint → Inference". Phía trước nó là ô 5A (tách token), phía sau là vòng decode. Phần tính toán bên trong là chuỗi ô 2A đến 2D (forward) với các trọng số đã nạp từ checkpoint, không có ô 3A đến 3E.

Vấn đề cần giải. Để sinh token đầu tiên, mô hình cần hiểu cả prompt. Ngoài ra các token sinh sau sẽ lại cần Key và Value của prompt; nếu không lưu lại, mỗi bước sau phải tính lại toàn bộ prompt. Prefill làm cả hai việc trong một lượt.

Khái niệm tối thiểu. Prompt có N token. Mỗi lớp có n_kv đầu key/value, mỗi đầu có d_head chiều. KV cache lưu K và V cho mọi vị trí đã xử lý, ở mọi lớp. Số số trong cache = 2 (K và V) × số lớp × n_kv × số vị trí × d_head.

Ví dụ có số (minh họa). Prompt [1, 3] (“mèo ăn”) nên N = 2. Mô hình thu nhỏ: 2 lớp, n_kv = 2, d_head = 2, kiểu số 32 bit (4 byte). Sau prefill, cache có 2 × 2 × 2 × 2 × 2 = 32 số tức 128 byte; mỗi lớp giữ K và V shape (2, 2, 2) = (số đầu, số vị trí, d_head). Mỗi token mới sau này thêm 2 × 2 × 2 × 2 = 16 số = 64 byte.

Logits. Forward cho logits ở cả hai vị trí, nhưng chỉ logits ở vị trí cuối (sau “ăn”) dùng để chọn token tiếp theo; logits vị trí đầu (dự đoán chữ sau “mèo”) chính là bài tập của lúc huấn luyện và ở đây bị bỏ qua. Giả sử sau huấn luyện, logits ở vị trí cuối là [0,1; 0,2; 0,1; 0,3; 2,4; 0,2; 0,5] cho [pad, mèo, chó, ăn, cá, ngủ, eos] (số minh họa). Greedy chọn chỉ số 4, tức “cá”. Softmax cho xác suất của cá ≈ 0,59, các token khác từ 0,06 đến 0,09.

Cách đọc đầu ra. Có hai kết quả: cache đã đầy cho N vị trí và token đầu tiên “cá” (ID 4). Lưu ý token vừa chọn chưa có Key và Value trong cache, nó sẽ được tính khi được đưa vào forward ở bước decode đầu tiên. Cũng lưu ý ‘đã chọn cá’ chỉ nghĩa là cá có điểm cao nhất ở checkpoint này; không phải câu trả lời đã đúng hay hữu ích, và với sampling khác (nhiệt độ cao) kết quả có thể khác.

Thực thi và cầu nối sang Inference Engineering. Engine có thể chia prompt dài thành từng đoạn (chunked prefill) hoặc bỏ qua phần prefix đã có cache hợp lệ; kernel attention giảm số lần đọc ghi bộ nhớ giúp xử lý context hiệu quả hơn. Prefill thường tốn nhiều tính toán nhưng điểm nghẽn thực sự phụ thuộc mô hình, độ dài và kích thước nhóm, nên phải đo trước khi quyết định tối ưu phép nhân ma trận hay bộ nhớ. Xem xử lý prompt trước khi sinh, bộ nhớ đệm key và value và chọn token từ điểm số để đi sâu từng phần; tách máy xử lý prompt khỏi máy sinh token mô tả một cách tổ chức khi hai pha có đặc tính khác nhau.

Để chạy prefill bằng trọng số đã học và kiểm tra token đầu được chọn bằng greedy, mở giai đoạn 13 · Từ checkpoint sang suy luận của Bài Lab tổng hợp. Hãy tự đoán trước kích thước cache sau prefill từ số lớp, số đầu, N và d_head.

4Profiling

• Time to First Token

  Thời gian từ lúc gửi yêu cầu đến token đầu tiên. Tách thành các phần: chờ trong hàng đợi, chuẩn bị đầu vào, prefill và xuất kết quả để biết phần nào chậm. Đo theo độ dài prompt khác nhau, không chỉ một độ dài, và ghi tải đồng thời.

• Tốc độ xử lý prompt và dùng lại prefix

  Đếm số token đầu vào thực sự phải tính (không tính phần đã có trong cache), và so lần đầu với lần dùng lại cùng prefix. Trước khi nói prefix caching giúp ích, kiểm tra kết quả logits của hai đường (có cache và không có cache) trùng trong sai số số học.

• Kích thước cache sau prefill

  Tính 2 × số lớp × n_kv × N × d_head × số byte mỗi số và đối chiếu với bộ nhớ thực tế bộ lập lịch đã đặt. Với ví dụ nhỏ là 128 byte. Sai khác lớn cho thấy cấu hình (số đầu, kiểu số) khác với giả định.

Mẹo đo lường: ghim số liệu trên workload thật, ghi cả p50/p95/p99 thay vì chỉ trung bình, và luôn ghi lại điều kiện chạy (batch, cache, đồng thời, phần cứng) trước và sau khi tối ưu. Số đo thiếu điều kiện chạy thì không so sánh được.

5Ví dụ

6Thuật ngữ

Prefill

Pha xử lý toàn bộ prompt đã biết, song song, trước khi sinh token mới.

KV cache

Bộ nhớ giữ cặp Key và Value của mọi token đã xử lý ở mỗi lớp để các bước sau không phải tính lại.

Logits ở vị trí cuối

Trong V điểm số của vị trí cuối prompt, điểm của mỗi token cho biết khả năng nó là token kế tiếp; sampling chọn token đầu tiên từ đây.

Time to First Token (TTFT)

Thời gian từ lúc gửi yêu cầu đến lúc có token đầu tiên, gồm chờ, tách token, prefill và chọn token.

Greedy decoding

Luôn chọn token có điểm cao nhất; xác định, không ngẫu nhiên.

7Bài liên quan

8Nguồn