Prefill và decode: KV cache và bức tường bộ nhớ

Vì sao token đầu tiên đến chậm còn các token sau đều đều, KV cache làm gì, và vì sao khi sinh chữ, thứ giới hạn tốc độ là băng thông bộ nhớ chứ không phải sức tính. Thêm KV cache vào GPT tự viết, đo, và thử số thực 16-bit.

Bài trước viết lại Qwen2.5-0.5B bằng 70 dòng PyTorch và để lại hai câu đố:

  • Vòng lặp sinh chữ của ta mất 55,7 giây cho 300 token, còn transformers chỉ mất 8,9 giây.
  • Đưa một token qua mô hình mất 28 ms, nhưng đưa 500 token cùng lúc chỉ mất 380 ms, tức 0,76 ms mỗi token. Số phép tính cho mỗi token thì như nhau.

Câu đố đầu có lời giải gọn là KV cache. Câu đố thứ hai sâu hơn, và lời giải của nó giải thích gần như mọi thứ về tốc độ của LLM: vì sao token đầu tiên đến chậm, vì sao API tính tiền token đầu ra đắt hơn token đầu vào, vì sao mô hình được nén xuống 8-bit hay 4-bit, và vì sao một server phục vụ nhiều người cùng lúc gần như không tốn thêm.

Lộ trình: hai giai đoạn của một câu trả lời, thêm KV cache vào mô hình tự viết, đo xem decode thực sự bị giới hạn bởi cái gì, một thí nghiệm quyết định với số thực 16-bit, KV cache tốn bao nhiêu bộ nhớ, chuyện gì xảy ra khi prompt dài, rồi các hệ quả thực tế.

Mọi con số trong bài đều đo trên máy thật, với Python 3.12, PyTorch 2.14, transformers 5.20, chạy bằng CPU. Code là bản tự viết ở bài trước, cộng thêm vài dòng.

Hai giai đoạn của một câu trả lời

Khi bạn gửi một câu hỏi, mô hình làm hai việc rất khác nhau:

  1. Prefill: đọc toàn bộ prompt. Mọi token của prompt đã biết sẵn, nên chúng đi qua mô hình cùng một lần, như một ma trận nhiều hàng. Kết thúc prefill là token đầu tiên của câu trả lời.
  2. Decode: sinh từng token còn lại. Token thứ n phải biết token thứ n − 1 là gì, nên không thể làm song song: mỗi bước một token.

Hai giai đoạn này cho ra hai con số mà mọi dịch vụ LLM đều theo dõi: thời gian tới token đầu tiên (time to first token, TTFT), do prefill quyết định, và số token mỗi giây khi chữ đang chạy, do decode quyết định. Chatbot ở seri Python có token đầu tiên sau 0,7 giây rồi chữ chạy đều khoảng 17 token mỗi giây: 0,7 giây đó là prefill.

KV cache: đừng tính lại những gì đã tính

Vòng lặp ở bài trước đưa cả dãy qua mô hình ở mỗi bước. Nhưng hãy nhìn lại attention: ở mỗi lớp, mỗi token tạo ra một key và một value, và chúng chỉ phụ thuộc vào token đó và các token trước nó. Khi thêm token thứ 301, key và value của 300 token trước không đổi một chút nào. Tính lại chúng là lãng phí thuần túy.

KV cache giữ lại key và value của mọi token đã qua, ở mọi lớp. Mỗi bước decode chỉ cần đưa một token mới qua mô hình; token đó tạo ra q, k, v của riêng nó, nối k và v vào cache, rồi nhìn vào toàn bộ cache. Sửa hàm attention của bài trước:

def attention(x, i, positions, cache):
    p = f"layers.{i}.self_attn."
    seq = x.shape[0]
    q = (x @ W[p + "q_proj.weight"].T + W[p + "q_proj.bias"]).view(seq, H, HD).transpose(0, 1)
    k = (x @ W[p + "k_proj.weight"].T + W[p + "k_proj.bias"]).view(seq, KV, HD).transpose(0, 1)
    v = (x @ W[p + "v_proj.weight"].T + W[p + "v_proj.bias"]).view(seq, KV, HD).transpose(0, 1)
    q, k = rope(q, positions), rope(k, positions)
    if i in cache:                                   # keys and values of every earlier token, computed once
        k = torch.cat([cache[i][0], k], dim=1)
        v = torch.cat([cache[i][1], v], dim=1)
    cache[i] = (k, v)
    k, v = k.repeat_interleave(H // KV, dim=0), v.repeat_interleave(H // KV, dim=0)
    scores = q @ k.transpose(-1, -2) / HD**0.5       # (14, new, all): the new tokens look at everything so far
    total = k.shape[1]
    causal = torch.ones(seq, total, dtype=torch.bool).triu(total - seq + 1)
    weights = scores.masked_fill(causal, float("-inf")).softmax(-1)
    return (weights @ v).transpose(0, 1).reshape(seq, D) @ W[p + "o_proj.weight"].T


def forward(ids, cache):
    """Only the new tokens go in; the cache holds what the model already knows about the earlier ones."""
    start = cache[0][0].shape[1] if 0 in cache else 0
    positions = torch.arange(start, start + len(ids))
    x = W["embed_tokens.weight"][ids]
    for i in range(LAYERS):
        x = x + attention(rms_norm(x, W[f"layers.{i}.input_layernorm.weight"]), i, positions, cache)
        x = x + mlp(rms_norm(x, W[f"layers.{i}.post_attention_layernorm.weight"]), i)
    return rms_norm(x, W["norm.weight"]) @ W["embed_tokens.weight"].T

Ba thay đổi so với bản cũ: k và v được nối vào cache, vị trí bắt đầu (start) lấy từ độ dài cache để RoPE xoay đúng góc, và mặt nạ causal bây giờ là hình chữ nhật (các token mới nhìn được mọi token cũ, cộng các token mới đứng trước chúng). Vòng lặp sinh chữ trở thành:

cache = {}
next_id = int(forward(prompt_ids, cache)[-1].argmax())  # prefill: the whole prompt at once
for _ in range(299):
    next_id = int(forward([next_id], cache)[-1].argmax())  # decode: one token at a time

Cùng câu hỏi ở bài trước, 300 token:

Sau token 1 100 200 300
Không cache 0,14 s 11,8 s 30,4 s 55,7 s
Có cache 0,08 s 3,0 s 6,0 s 9,0 s
transformers generate() 8,9 s

300 token sinh ra giống hệt từng token với bản không cache, và thời gian bây giờ tăng đều theo đường thẳng, khoảng 30 ms mỗi token, ngang với transformers. Câu đố thứ nhất đã xong. Nhưng chính con số 30 ms đó mới là câu đố thật.

Decode bị giới hạn bởi cái gì

Một bước decode đưa một token qua mô hình. Hãy đếm xem nó phải làm gì:

  • Tính toán: mỗi tham số tham gia đúng một phép nhân và một phép cộng, nên khoảng 2 × 494 triệu ≈ 1 tỉ phép tính.
  • Đọc bộ nhớ: mỗi tham số phải được đọc từ RAM ít nhất một lần. 494 triệu số thực 32-bit là 1,98 GB.

Đo trên cùng máy, ở hai chế độ:

Thời gian Phép tính mỗi giây Byte trọng số đọc mỗi giây
Decode: 1 token 27,6 ms 36 tỉ 72 GB/s
Prefill: 1.000 token cùng lúc 634 ms 1.600 tỉ 3 GB/s

Để biết 72 GB/s là nhiều hay ít, đo thứ đơn giản nhất có thể: cộng hết một mảng 2 GB (x.sum()). Phép đó gần như không tính gì, chỉ đọc. Kết quả: 83 GB/s. Đó là tốc độ đọc RAM trần của máy này.

Vậy là rõ: khi decode, mô hình đọc trọng số ở 72 GB/s, tức 87% tốc độ đọc RAM tối đa, trong khi phép tính chỉ dùng chưa tới 1/40 sức tính mà chính máy đó đạt được khi prefill. Decode không bị giới hạn bởi sức tính; nó bị giới hạn bởi tốc độ kéo trọng số từ bộ nhớ vào bộ xử lý. Bộ xử lý phần lớn thời gian ngồi chờ dữ liệu.

Prefill thì ngược lại. Trọng số được đọc một lần nhưng dùng cho 1.000 token, nên mỗi byte đọc vào được dùng cho hàng nghìn phép tính. Ở đó thứ giới hạn là sức tính. Con số hay dùng để phân biệt hai chế độ là số phép tính trên mỗi byte đọc vào: decode ở số thực 32-bit làm khoảng 2 phép tính trên 4 byte, tức 0,5; prefill 1.000 token làm khoảng 500. Phần cứng nào cũng có một ngưỡng ở giữa: dưới ngưỡng, bạn chờ bộ nhớ; trên ngưỡng, bạn chờ bộ xử lý.

Đây chính là lời giải của câu đố thứ hai. 28 ms cho một token không phải là thời gian tính; đó là thời gian đọc 2 GB. Đọc 2 GB một lần rồi dùng cho 500 token thì mỗi token chỉ chịu 1/500 cái giá đó.

Thí nghiệm quyết định: số thực 16-bit

Nếu giả thuyết "decode chờ bộ nhớ" là đúng, thì giảm một nửa số byte của trọng số phải làm decode nhanh lên gần gấp đôi, dù số phép tính không đổi. Và nếu "prefill chờ sức tính" là đúng, prefill sẽ không được lợi gì. Chạy cùng mô hình ở ba kiểu số:

Kiểu số Trọng số 1 token (decode) 500 token (prefill)
float32 1,98 GB 27,3 ms 341 ms
bfloat16 0,99 GB 16,5 ms 3.711 ms
float16 0,99 GB 16,3 ms 2.501 ms

Cả hai dự đoán đều đúng, theo cách rõ hơn cả mong đợi:

  • Decode nhanh lên 1,7 lần chỉ nhờ trọng số nhỏ đi một nửa. Không có phép tính nào được bỏ bớt; chỉ là có ít byte hơn phải kéo qua.
  • Prefill chậm đi 7–11 lần. Trên CPU này, phép nhân ma trận 16-bit chậm hơn nhiều so với 32-bit, và vì prefill bị giới hạn bởi sức tính, nó chịu đủ cái chậm đó. Decode cũng làm cùng những phép nhân 16-bit chậm ấy, nhưng vì phần lớn thời gian nó đang chờ bộ nhớ, cái chậm đó gần như bị che khuất (tốc độ đọc trọng số chỉ tụt từ 72 xuống 60 GB/s).

Cùng một thay đổi, hai kết quả ngược nhau, vì hai giai đoạn bị giới hạn bởi hai thứ khác nhau. GPU có phần cứng riêng cho phép nhân 16-bit nên prefill không bị chậm như thế, nhưng logic thì giống hệt. Đây là lý do mô hình được lượng tử hóa (quantize) xuống 8-bit hay 4-bit khi chạy: với decode, mỗi bit bỏ đi là tốc độ thêm vào. Bài về phục vụ nhiều người sẽ đo chuyện đó.

Decode chậm dần theo độ dài

Có cache, mỗi bước decode vẫn không hoàn toàn bằng nhau. Một bước decode sau các prompt dài ngắn khác nhau:

Ngữ cảnh (token) Một bước decode
10 27,6 ms
100 29,0 ms
1.000 33,9 ms
4.000 51,0 ms
8.000 72,1 ms

Ngoài 1,98 GB trọng số, mỗi bước còn phải đọc toàn bộ KV cache: token mới nhìn vào key và value của mọi token trước nó. Ngữ cảnh càng dài, cache càng lớn, và ở 8.000 token, bước decode chậm gần gấp ba. Với mô hình lớn và ngữ cảnh rất dài, đọc KV cache có thể tốn ngang hoặc hơn đọc trọng số.

KV cache tốn bao nhiêu bộ nhớ

Mỗi token, ở mỗi lớp, giữ một key và một value cho mỗi head key/value:

byte mỗi token = số lớp × 2 (k và v) × số head KV × kích thước head × byte mỗi số

Tính từ file cấu hình của ba mô hình Qwen2.5, ở bfloat16:

Mô hình Lớp Head query / KV KV cache mỗi token Nếu không có grouped-query 32.768 token
0,5B 24 14 / 2 12 KiB 84 KiB 0,38 GiB
7B 28 28 / 4 56 KiB 392 KiB 1,75 GiB
72B 80 64 / 8 320 KiB 2.560 KiB 10 GiB

Một cuộc hội thoại 32.000 token với mô hình 72B giữ 10 GiB cache, chỉ cho một người dùng. Không có grouped-query attention (bài trước), con số đó là 80 GiB, gấp 8 lần. Đây là lý do thật sự của việc 14 head query chỉ dùng chung 2 head key/value: không phải để tính nhanh hơn, mà để cache nhỏ đi đủ để vừa bộ nhớ và đọc nhanh hơn ở mỗi bước decode.

Prompt dài: ma trận attention bùng nổ

Prefill có vấn đề riêng của nó. Đo thời gian prefill theo độ dài prompt, với hàm attention tự viết:

Prompt Prefill Mỗi token Ma trận scores mỗi lớp
1.000 843 ms 0,84 ms 0,06 GB
4.000 7.656 ms 1,91 ms 0,90 GB
8.000 87.942 ms 10,99 ms 3,58 GB

Từ 4.000 lên 8.000 token, thời gian không gấp đôi mà gấp 11 lần. Thủ phạm là dòng scores = q @ k.transpose(-1, -2): nó tạo ra một ma trận 14 × seq × seq, và ở 8.000 token, ma trận đó là 3,58 GB, cộng thêm mặt nạ và kết quả softmax cùng cỡ, ở mỗi lớp. Máy phải xoay xở với hàng chục GB dữ liệu tạm chỉ để rồi nhân chúng với v và vứt đi.

Thay đúng phần đó bằng hàm dựng sẵn của PyTorch:

out = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)
Prompt Attention tự viết scaled_dot_product_attention
1.000 843 ms 634 ms
4.000 7.656 ms 3.162 ms
8.000 87.942 ms 8.709 ms

Kết quả giống nhau (chênh lệch logit lớn nhất 0,0001), nhưng ở 8.000 token nhanh hơn 10 lần. Hàm này không bao giờ tạo ra cả ma trận: nó chia q, k, v thành từng khối nhỏ, tính softmax từng phần, và cộng dồn kết quả. Đây là ý tưởng của FlashAttention: phép tính vẫn như cũ, nhưng dữ liệu tạm luôn nằm gọn trong bộ nhớ nhanh. Lại là cùng một bài học: thứ đắt là di chuyển dữ liệu, không phải tính toán.

Phép tính của attention vẫn tăng theo bình phương độ dài, nên từ 4.000 lên 8.000 token vẫn chậm hơn gấp đôi. Đó là lý do prompt dài vẫn đắt, và token đầu tiên của một câu hỏi kèm tài liệu 100 trang đến chậm hơn hẳn câu hỏi trống.

Hệ quả thực tế

Token đầu ra đắt hơn token đầu vào. Hầu hết API tính tiền token đầu ra đắt hơn token đầu vào nhiều lần. Bây giờ ta biết vì sao: token đầu vào đi qua prefill, cùng lúc với nhau, mỗi lần đọc trọng số chia cho hàng nghìn token; token đầu ra đi qua decode, mỗi token một lần đọc toàn bộ mô hình.

Prompt caching. Nhiều API có "prompt caching": nếu nhiều yêu cầu bắt đầu bằng cùng một đoạn (lời dặn hệ thống dài, tài liệu tham khảo), phần đó được tính rẻ hơn và token đầu tiên đến nhanh hơn. Đó chính là KV cache ở trên, chỉ là được giữ lại giữa các yêu cầu: key và value của đoạn đầu đã tính rồi thì không cần prefill lại. Muốn tận dụng nó, hãy đặt phần cố định ở đầu prompt và phần thay đổi (câu hỏi của người dùng) ở cuối.

Ước lượng tốc độ của một mô hình bất kỳ. Vì decode bị giới hạn bởi bộ nhớ, có một công thức nhanh cho tốc độ tối đa:

token mỗi giây (tối đa) ≈ băng thông bộ nhớ ÷ số byte trọng số

Kiểm tra với chính mô hình ở trên: 83 GB/s ÷ 1,98 GB ≈ 42 token/giây tối đa; thực tế đo được 36 (27,6 ms mỗi bước). Áp dụng cho các mô hình lớn hơn (đây là ước tính, không phải đo): Qwen2.5-7B ở bfloat16 nặng khoảng 15 GB, nên trên máy này chạy được nhiều nhất khoảng 5 token mỗi giây. Một GPU trung tâm dữ liệu với băng thông bộ nhớ khoảng 3.000 GB/s thì được khoảng 200. Khi chọn phần cứng để chạy LLM, con số cần nhìn đầu tiên không phải là sức tính mà là băng thông bộ nhớ và dung lượng bộ nhớ (để chứa trọng số và KV cache).

Phục vụ nhiều người gần như miễn phí. Câu đố thứ hai còn một mặt nữa. Nếu 8 người cùng chờ token kế tiếp, server có thể đưa 8 token (một token mỗi người) qua mô hình cùng lúc: trọng số được đọc một lần cho cả 8. Bài trước đã đo: 1 token mất 28 ms, 10 token cùng lúc mất 52 ms. Phục vụ 10 người tốn chưa tới gấp đôi phục vụ 1 người. Đó là nền tảng của batching, chủ đề của bài phục vụ nhiều người.

Tóm lại

Khái niệm Thực chất
Prefill Cả prompt đi qua mô hình một lần; quyết định thời gian tới token đầu tiên; giới hạn bởi sức tính
Decode Mỗi bước một token; quyết định token mỗi giây; giới hạn bởi băng thông bộ nhớ
KV cache Giữ key và value của mọi token đã qua; mỗi bước decode chỉ tính một token mới
Bức tường bộ nhớ Decode đọc trọng số ở 87% tốc độ RAM tối đa, chỉ dùng chưa tới 1/40 sức tính
16-bit Trọng số nhỏ một nửa: decode nhanh 1,7 lần; prefill trên CPU chậm đi nhiều lần
Kích thước cache Lớp × 2 × head KV × kích thước head × byte; 10 GiB cho 32K token ở mô hình 72B
Grouped-query attention Cache nhỏ đi 7–8 lần; lợi ích nằm ở bộ nhớ, không ở phép tính
FlashAttention Không bao giờ tạo cả ma trận attention; 10 lần nhanh hơn ở 8.000 token
Ước lượng tốc độ Token/giây tối đa ≈ băng thông bộ nhớ ÷ byte trọng số

Ba bài đầu đã đi hết đường đi của một câu trả lời: chữ thành token, token qua mô hình, và vì sao đường đi đó nhanh hay chậm. Nhưng 494 triệu con số trong file trọng số từ đâu ra? Bài sau sẽ huấn luyện một mô hình từ đầu, trên văn bản tiếng Việt, ngay trên CPU, để xem những con số ngẫu nhiên dần trở thành thứ biết viết chữ như thế nào.