Inference ở quy mô lớn: KV cache, prefill và decode

Bài thứ mười hai kết thúc bằng một lời hứa: inference chia thành hai pha rất khác nhau, prefill bị giới hạn bởi compute còn decode bị giới hạn bởi băng thông bộ nhớ. Bài này giải thích vì sao, và đi từ đó tới mọi kỹ thuật mà một inference server hiện đại dùng.

Câu hỏi của bài: vì sao cùng một GPU, phục vụ một người và phục vụ một trăm người lại khác nhau đến vậy? Một người dùng chat với model 8B trên GPU riêng thấy vài chục token mỗi giây, và GPU gần như đứng yên phần lớn thời gian. Một trăm người dùng cùng lúc thì tổng throughput tăng gần trăm lần, cho tới khi bộ nhớ hết chỗ. Hiểu được hai câu đó là hiểu được KV cache, batching, PagedAttention, FlashAttention và speculative decoding.

Toàn bài xoay quanh một câu hỏi duy nhất: mỗi byte đọc từ bộ nhớ làm được bao nhiêu việc? Mọi kỹ thuật phía sau hoặc đọc ít byte hơn cho mỗi token, hoặc làm nhiều việc hơn cho mỗi byte đã đọc. Bài đi theo thứ tự của Hình 1 ở đầu trang.

01Từ một model đến một dịch vụ

Bài 2, mục 8 đã vẽ bức tranh chung: hai metric latency và throughput thường xung đột (mục 8.1), generation là autoregressive nên không song song theo thời gian được (8.2), KV cache đổi VRAM lấy compute (8.3), PagedAttention quản lý cache như virtual memory (8.4), continuous batching chèn request mới khi có slot trống (8.5), và FlashAttention tối ưu ở tầng kernel (8.6). Bài này không lặp lại các phần đó mà đi vào câu hỏi bài 2 bỏ ngỏ: vì sao những kỹ thuật ấy hiệu quả, và hiệu quả tới đâu.

Ba metric dùng xuyên suốt bài:

Người dùng chat cảm nhận TTFT (chờ bao lâu mới thấy chữ) và TPOT (chữ chạy nhanh hay chậm). Người trả tiền GPU quan tâm throughput. Phần còn lại của bài giải thích vì sao ba con số này phụ thuộc vào những thứ rất khác nhau.

02Prefill và decode: hai pha, hai nút thắt

Hai pha của một request

Khi một request đến, model chạy hai pha:

• Prefill. Cả prompt n token đi qua model trong một lượt forward. Mọi phép nhân với ma trận trọng số là ma trận × ma trận: n vector đầu vào nhân với cùng một ma trận weights. K và V của mọi token prompt được ghi vào KV cache. Kết quả là logits của vị trí cuối, cho ra token đầu tiên. Prefill quyết định TTFT.
• Decode. Mỗi bước chỉ có một token mới (cho mỗi sequence). Phép nhân với weights là ma trận × vector. Token mới tính Q, K, V của riêng nó, ghi thêm một dòng vào KV cache, và attention đọc lại toàn bộ K, V cũ. Bước này lặp lại cho tới EOS. Decode quyết định TPOT.

Cả hai pha làm cùng một loại phép tính, nhưng tỉ lệ giữa tính và đọc thì khác hẳn.

Đếm FLOPs và đếm byte

Bài 8, mục 6 đã đếm: một lượt forward tốn khoảng 2N FLOPs cho mỗi token, với N là số tham số, cộng một số hạng attention tăng theo độ dài context. Bài này chỉ cần thêm một phép đếm nữa: mỗi lượt forward phải đọc toàn bộ weights từ bộ nhớ GPU (HBM, hay GDDR6 trên card consumer) vào chip để nhân. Ở fp16 là 2 byte mỗi tham số, tức 2N byte.

Tỉ số giữa hai con số đó gọi là arithmetic intensity: số FLOP làm được cho mỗi byte đọc vào.

• Decode, batch 1: 2N FLOPs trên 2N byte weights, tức khoảng 1 FLOP/byte.
• Decode, batch B: một lần đọc weights phục vụ B sequence, nên 2N·B FLOPs trên 2N byte, khoảng B FLOP/byte (chưa tính KV cache, phần sau sẽ tính).
• Prefill, prompt n token: một lần đọc weights phục vụ n token, khoảng n FLOP/byte.

Roofline và ridge point của RTX 2080 Ti

Mô hình roofline (Williams, Waterman & Patterson 2009) nói rằng một kernel bị giới hạn bởi cái chậm hơn trong hai thứ: tốc độ tính, hoặc tốc độ đọc dữ liệu. Thời gian tối thiểu là:

t ≥ max( FLOPs / peak_FLOPS ,  bytes / bandwidth )

Điểm chuyển giữa hai chế độ gọi là ridge point: peak_FLOPS / bandwidth. Kernel có arithmetic intensity thấp hơn ridge point là memory-bound (GPU ngồi chờ dữ liệu), cao hơn là compute-bound.

Lấy RTX 2080 Ti, card dùng ở các bài trước. Theo Bảng 1 của NVIDIA Turing Architecture Whitepaper (bản reference, không phải Founders Edition):

Khi nhân ma trận fp16, PyTorch thường cộng dồn ở fp32 để giữ độ chính xác, nên bài lấy 53,8 TFLOPS làm mốc so sánh. Ridge point:

53,8e12 FLOP/s / 616e9 byte/s ≈ 87 FLOP/byte     (cộng dồn FP32)
107,6e12 / 616e9              ≈ 175 FLOP/byte    (cộng dồn FP16)

Để so sánh, kipply tính cho A100: 312 TFLOPS bf16 và 1,5 TB/s, tỉ lệ khoảng 208 (bài Transformer Inference Arithmetic, mục “kv cache”). GPU càng mới thì compute tăng nhanh hơn băng thông, ridge point càng cao, và decode càng “lãng phí” compute.

Decode ở batch 1 có intensity khoảng 1, thấp hơn ridge point gần 90 lần. Nghĩa là trong lúc decode một sequence, Tensor Core của 2080 Ti nhàn rỗi khoảng 99% thời gian. Đó là lý do câu “decode bị giới hạn bởi băng thông” của bài 12.

Bảng FLOP/byte cho Qwen2.5-1.5B

Áp vào một model vừa với 2080 Ti: Qwen2.5-1.5B ở fp16 (1,54B tham số, 3,09 GB weights; 28 layer, 12 query head, 2 KV head, head_dim 128, tức 28 KiB KV cache mỗi token, cách tính ở mục 3). Bảng dưới là phép tính lý thuyết theo roofline, không phải số đo: FLOPs gồm 2N cộng phần attention, byte gồm weights cộng KV cache đọc vào.

Hai điều đọc được từ bảng:

1. Decode gần như “miễn phí” khi tăng batch, cho tới một ngưỡng. Từ batch 1 lên batch 16, thời gian tối thiểu mỗi bước chỉ tăng từ 5,1 lên 5,8 ms, trong khi số token sinh ra mỗi bước tăng 16 lần. Phần lớn thời gian là đọc 3,09 GB weights, và lần đọc đó được chia cho cả batch.
2. Ngưỡng không chỉ là compute mà còn là KV cache. Ở batch 256 với context 1.024, KV cache đã là 256 × 1.024 × 28 KiB ≈ 7,5 GB, cộng 3,09 GB weights là sát 11 GB của card. Trước khi chạm ridge point, ta đã hết bộ nhớ.

Vì sao KV cache kéo intensity xuống

Weights được chia cho cả batch, còn KV cache thì không: mỗi sequence có K, V riêng, và mỗi bước decode phải đọc lại toàn bộ. Phần attention của decode có intensity cố định, không tăng theo batch: mỗi phần tử K đọc vào (2 byte) được g query head dùng chung trong GQA, mỗi head một phép nhân–cộng (2 FLOP), nên intensity của attention khi decode xấp xỉ g FLOP/byte, với g là số query head mỗi KV head. Qwen2.5-1.5B có g = 6, Llama 3.1 8B có g = 4, model MHA có g = 1.

Vì vậy khi context dài, decode càng tụt sâu vào vùng memory-bound dù batch lớn. Cùng batch 64 trên Qwen2.5-1.5B, intensity lý thuyết giảm từ khoảng 42 ở context 1.024 xuống khoảng 23 ở context 4.096 và khoảng 9 ở context 32.768 (lúc đó KV cache đã là 60 GB, không vừa card nào trong tầm tay). Đây là sợi chỉ nối mọi kỹ thuật phía sau: bộ nhớ dành cho KV cache quyết định batch tối đa, và batch quyết định throughput.

Vì sao weight-only quantization giúp decode mà không giúp prefill

Bài 12, mục 5–6 đã dẫn số đo trong README của llama.cpp: ở Q4_K_M, text generation nhanh hơn F16 khoảng 2,5 lần (29 lên 72 token/s), còn prompt processing chậm đi một chút (923 xuống 822 token/s). Roofline giải thích ngay: decode memory-bound, nên weights nhỏ đi 4 lần là byte đọc vào giảm gần 4 lần; prefill compute-bound, phép nhân vẫn chạy ở 16 bit sau khi dequantize, nên không nhanh hơn, và bước dequantize còn làm nó chậm đi một chút.

Code mẫu 1: đo TTFT và thời gian mỗi bước decode

Code dưới đo (a) thời gian prefill theo độ dài prompt và (b) thời gian mỗi bước decode theo batch, rồi quy ra băng thông đạt được để so với 616 GB/s.

# Code minh hoạ. Python 3.12, torch 2.14.1+cu130, transformers 5.18.0
import time, torch
from transformers import AutoModelForCausalLM

name = "Qwen/Qwen2.5-1.5B"
model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float16, device_map="cuda").eval()
cfg = model.config
head_dim = cfg.hidden_size // cfg.num_attention_heads
kv_per_token = 2 * cfg.num_hidden_layers * cfg.num_key_value_heads * head_dim * 2  # byte, fp16
weight_bytes = sum(p.numel() * p.element_size() for p in model.parameters())

def sync_time():
    torch.cuda.synchronize()
    return time.perf_counter()

@torch.no_grad()
def prefill(ids):
    # logits_to_keep=1: chỉ tính logits ở vị trí cuối, tránh tensor batch × n × vocab
    return model(ids, use_cache=True, logits_to_keep=1)

# (a) Prefill ≈ TTFT khi không phải xếp hàng
for n in (128, 512, 2048, 8192):
    ids = torch.randint(1000, 30000, (1, n), device="cuda")
    prefill(ids)                                    # warmup
    t0 = sync_time(); prefill(ids); dt = sync_time() - t0
    print(f"prompt {n:5d}: prefill {dt * 1e3:7.1f} ms  ({n / dt:8.0f} token/s)")

# (b) Mỗi bước decode theo batch, context ban đầu 1.024 token
ctx, steps = 1024, 32
for B in (1, 4, 16, 64):
    ids = torch.randint(1000, 30000, (B, ctx), device="cuda")
    with torch.no_grad():
        out = prefill(ids)
        past, nxt = out.past_key_values, out.logits[:, -1:].argmax(-1)
        t0 = sync_time()
        for _ in range(steps):
            out = model(nxt, past_key_values=past, use_cache=True)
            past, nxt = out.past_key_values, out.logits[:, -1:].argmax(-1)
        dt = (sync_time() - t0) / steps
    kv_read = B * (ctx + steps // 2) * kv_per_token  # KV trung bình đọc mỗi bước
    print(f"batch {B:3d}: {dt * 1e3:6.2f} ms/bước  {B / dt:8.0f} token/s  "
          f"≈ {(weight_bytes + kv_read) / dt / 1e9:5.0f} GB/s")

Chưa chạy thử. Cần GPU NVIDIA còn trống khoảng 8 GB VRAM trở lên (ví dụ RTX 2080 Ti 11 GB), CUDA 13.0, Python 3.12, torch 2.14.1+cu130, transformers 5.18.0.

Điều nên thấy khi chạy: ở phần (b), ms mỗi bước tăng rất chậm từ batch 1 tới batch 16, nên token/s tăng gần tuyến tính; băng thông quy đổi tiến dần về phía 616 GB/s nhưng không vượt được. Ở phần (a), token/s của prefill cao hơn decode batch 1 hàng trăm lần. Bảng lý thuyết ở trên cho biết nên kỳ vọng thứ tự độ lớn nào; con số thật luôn kém hơn giới hạn dưới vì overhead của kernel và Python.

03KV cache: toán bộ nhớ

Byte mỗi token

Công thức kích thước KV cache đã có ở bài 2, mục 8.3. Viết lại theo byte cho mỗi token thì dễ so giữa các model hơn:

KV byte mỗi token = 2 × num_layers × num_kv_heads × head_dim × byte_mỗi_giá_trị

Số 2 là cho K và V. Nhân con số này với tổng số token đang nằm trong batch (mọi sequence cộng lại) là ra tổng KV cache. Paper vLLM lấy ví dụ OPT-13B (MHA, 40 layer, hidden 5.120): 800 KB mỗi token ở fp16, và một request 2.048 token có thể tốn tới 1,6 GB (Kwon et al. 2023, mục 3, đoạn “Large KV cache”).

KV cache chỉ lưu K và V, không lưu Q và không lưu attention score. Q của token cũ không bao giờ được dùng lại: ở bước decode, chỉ Q của token mới cần nhân với mọi K cũ (bài 7, mục 7 giải thích vì sao causal mask làm token cũ không bao giờ phải “nhìn lại” token mới).

Code mẫu 2: tính KV cache từ config.json

Thay vì tra bảng, đọc thẳng config.json trên Hugging Face Hub. Code dưới cũng đếm số tham số từ header của các file safetensors, không cần tải weights.

# Python 3.12, huggingface_hub 1.33.0 — chỉ tải config.json và header safetensors, chạy trên CPU
import json, os
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
from huggingface_hub import hf_hub_download, get_safetensors_metadata

MODELS = {
    "Llama 3.1 8B":   "unsloth/Meta-Llama-3.1-8B",  # bản mirror không gated, config giống hệt meta-llama
    "Llama 2 7B":     "NousResearch/Llama-2-7b-hf", # MHA: số KV head bằng số query head
    "Qwen2.5-7B":     "Qwen/Qwen2.5-7B",
    "Mistral 7B v0.3": "mistralai/Mistral-7B-v0.3",
    "DeepSeek-V3":    "deepseek-ai/DeepSeek-V3",    # MLA
}
KiB, GiB = 1024, 1024**3

def kv_bytes_per_token(cfg, bytes_per_value=2):
    L = cfg["num_hidden_layers"]
    if "kv_lora_rank" in cfg:                      # MLA: cache latent c_KV và phần RoPE của key
        return L * (cfg["kv_lora_rank"] + cfg["qk_rope_head_dim"]) * bytes_per_value
    n_kv = cfg.get("num_key_value_heads", cfg["num_attention_heads"])
    head_dim = cfg.get("head_dim") or cfg["hidden_size"] // cfg["num_attention_heads"]
    return 2 * L * n_kv * head_dim * bytes_per_value  # 2 = K và V

rows = {}
print(f"{'model':16} {'attn':4} {'L':>3} {'q/kv':>7} {'KiB/token':>9} {'GiB @128k':>9} {'weights GiB':>11}")
for name, repo in MODELS.items():
    cfg = json.load(open(hf_hub_download(repo, "config.json")))
    n_params = get_safetensors_metadata(repo).parameter_count  # đếm từ header, không tải weights
    per_tok = kv_bytes_per_token(cfg)
    n_q = cfg["num_attention_heads"]
    n_kv = cfg.get("num_key_value_heads", n_q)
    kind = "MLA" if "kv_lora_rank" in cfg else ("MHA" if n_kv == n_q else "GQA")
    weights = sum(n_params.values()) * 2             # giả định bf16 cho mọi tham số
    rows[name] = (cfg, per_tok, weights)
    print(f"{name:16} {kind:4} {cfg['num_hidden_layers']:>3} {f'{n_q}/{n_kv}':>7} "
          f"{per_tok / KiB:>9.1f} {per_tok * 131072 / GiB:>9.2f} {weights / GiB:>11.1f}")

# Ngân sách một GPU 80 GB cho Llama 3.1 8B: phần còn lại sau weights dành hết cho KV cache
cfg, per_tok, weights = rows["Llama 3.1 8B"]
budget = 80e9 * 0.9 - weights                       # chừa 10% cho activation, CUDA context...
n_params = weights / 2
non_embed = n_params - 2 * cfg["vocab_size"] * cfg["hidden_size"]  # bỏ embedding và lm_head
d_attn = cfg["num_attention_heads"] * cfg["head_dim"]
print(f"\nLlama 3.1 8B: weights {weights/1e9:.1f} GB, KV budget {budget/1e9:.1f} GB, N không embedding {non_embed/1e9:.2f}B")
print(f"{'context':>8} {'KV/seq GB':>9} {'max seqs':>8} {'prefill TFLOPs':>14} {'attention %':>11}")
for n in (4096, 32768, 131072):
    kv_seq = per_tok * n
    linear = 2 * non_embed * n                               # 2N mỗi token
    attn = 2 * cfg["num_hidden_layers"] * n * d_attn * n     # số hạng context của bài 8, n_ctx = n
    print(f"{n:>8} {kv_seq/1e9:>9.2f} {int(budget // kv_seq):>8} {(linear+attn)/1e12:>14.1f} {attn/(linear+attn):>11.0%}")
print(f"điểm cắt: attention = 2N khi n ≈ {non_embed / (cfg['num_hidden_layers'] * d_attn):,.0f} token")

Output (chạy ngày 08/10/2026):

model            attn   L    q/kv KiB/token GiB @128k weights GiB
Llama 3.1 8B     GQA   32    32/8     128.0     16.00        15.0
Llama 2 7B       MHA   32   32/32     512.0     64.00        12.6
Qwen2.5-7B       GQA   28    28/4      56.0      7.00        14.2
Mistral 7B v0.3  GQA   32    32/8     128.0     16.00        13.5
DeepSeek-V3      MLA   61 128/128      68.6      8.58      1275.0

Llama 3.1 8B: weights 16.1 GB, KV budget 55.9 GB, N không embedding 6.98B
 context KV/seq GB max seqs prefill TFLOPs attention %
    4096      0.54      104           61.6          7%
   32768      4.29       13          738.9         38%
  131072     17.18        3         6333.3         71%
điểm cắt: attention = 2N khi n ≈ 53,250 token

Llama 3.1 8B và Llama 2 7B có repo gốc gated trên Hugging Face, nên code dùng bản mirror (unsloth/…, NousResearch/…). Cấu hình khớp với Bảng 3 của paper Llama 3: 32 layer, 32 attention head, 8 key/value head cho bản 8B (Llama Team 2024, mục 3.2).

Đọc bảng

• Llama 3.1 8B tốn 128 KiB mỗi token. Một sequence 128k token tốn 16 GiB KV cache, xấp xỉ bằng cả phần weights bf16 (15 GiB). Ở context dài, một request “nặng” ngang một bản model.
• MHA đắt gấp 4 lần. Llama 2 7B với 32 KV head tốn 512 KiB mỗi token. Cùng ngân sách bộ nhớ, nó phục vụ được ít sequence hơn Llama 3.1 8B bốn lần, dù model nhỏ hơn.
• Qwen2.5-7B chỉ 56 KiB nhờ 4 KV head trên 28 query head, tức mỗi KV head phục vụ 7 query head.
• DeepSeek-V3 có 671B tham số nhưng KV cache chỉ 68,6 KiB mỗi token, ít hơn Llama 3.1 8B. Đó là nhờ MLA, mục 4 giải thích. Cột weights 1.275 GiB tính cho bf16; bản phát hành dùng FP8 nên nhỏ khoảng một nửa.

Phần dưới của output là ngân sách cho một GPU 80 GB: sau 16,1 GB weights và 10% dự phòng, còn khoảng 56 GB cho KV cache. Ở context 4k, GPU giữ được 104 sequence cùng lúc; ở 32k còn 13; ở 128k chỉ còn 3. Throughput của decode tỉ lệ với batch, nên cùng một GPU, context dài hơn 32 lần thì số request phục vụ đồng thời giảm khoảng 35 lần. Hai cột cuối dùng ở mục 9.

Ngân sách bộ nhớ của một GPU

Bộ nhớ GPU khi serving chia thành ba phần:

VRAM = weights + KV cache + (activation tạm thời, CUDA context, buffer của kernel)

Weights cố định. Activation ở inference nhỏ vì không cần giữ cho backward, chỉ giữ cho layer đang chạy. Phần còn lại dành cho KV cache, và chính phần đó quyết định số request đồng thời tối đa. Các inference server như vLLM cấp phát trước một tỉ lệ cố định của VRAM (tham số gpu_memory_utilization) rồi chia phần dư sau weights thành các block KV.

KV cache quantization

Đã lượng tử weights ở bài 12 thì cũng lượng tử được KV cache. Cách đơn giản nhất là lưu K, V ở FP8 thay vì fp16/bf16: mỗi giá trị 1 byte thay vì 2, tức KV cache nhỏ đi một nửa theo phép tính trực tiếp. Tài liệu vLLM ghi rằng lượng tử KV cache sang FP8 “giảm đáng kể bộ nhớ”, cho phép lưu nhiều token hơn, nên tăng throughput và hỗ trợ context dài hơn; hai định dạng là fp8_e4m3 và fp8_e5m2. Đi xa hơn, KIVI (Liu et al. 2024) lượng tử KV cache xuống 2 bit, với một quan sát giống chuyện outlier ở bài 12: key nên lượng tử theo kênh còn value theo token (abstract). Trên Llama-2-7B với một A100 80 GB, KIVI giảm 2,6 lần bộ nhớ đỉnh (kể cả weights), cho batch lớn hơn tới 4 lần và throughput cao hơn 2,35–3,47 lần (abstract, mục 4.2.4, Hình 5). Cái giá cũng giống bài 12: phải đo lại chất lượng trên tác vụ của mình. Hugging Face Transformers có sẵn QuantizedCache (cache_implementation="quantized") để thử.

04MQA, GQA và MLA

Nhìn vào công thức byte mỗi token, chỉ có hai núm vặn mà không đụng tới độ sâu model: số KV head và byte mỗi giá trị. Mục 3 vừa nói về núm thứ hai. Mục này nói về núm thứ nhất. Bài 8, mục 6 đã nói GQA giảm tham số của W_K, W_V và giảm KV cache nhưng không giảm FLOPs của QKᵀ; bài 1 mục 5.10 có hình minh hoạ MHA, MQA và GQA. Ở đây chỉ bàn chuyện bộ nhớ và chất lượng.

MQA: một cặp K/V cho mọi head

Multi-query attention (Shazeer 2019) giữ nhiều query head nhưng cho tất cả dùng chung một cặp K, V. Abstract viết thẳng mục tiêu: giảm kích thước các tensor K, V “và do đó giảm yêu cầu băng thông bộ nhớ của incremental decoding”. Trên WMT14 Anh–Đức, phần decoder giảm từ 46 xuống 3,8 micro giây cho mỗi token (Bảng 2, mục 4.3; đo trên TPUv2, đã chia đều cho batch 1.024), còn BLEU gần như giữ nguyên: 27,7 so với 27,5 khi greedy (Bảng 1, mục 4.2). Pope et al. 2022 cho thấy ở quy mô PaLM 540B, multiquery attention cho phép context dài hơn tới 32 lần trong cùng ngân sách bộ nhớ (abstract, Bảng 1 mục 4.2).

Nhược điểm: chất lượng có giảm, và train ổn định khó hơn. Với KV cache nhỏ đi n_heads lần, đó là đánh đổi lớn.

GQA: chia head thành nhóm

Grouped-query attention (Ainslie et al. 2023) đứng giữa: chia n_heads query head thành n_kv nhóm, mỗi nhóm dùng chung một cặp K, V. n_kv = n_heads là MHA, n_kv = 1 là MQA. Kết quả chính ở Bảng 1 (mục 3.2) của paper, trên T5-XXL:

GQA với 8 nhóm có chất lượng gần MHA và tốc độ gần MQA. Một điểm thực dụng: không cần train lại từ đầu. Paper uptrain checkpoint MHA bằng cách gộp các head K, V trong cùng nhóm bằng trung bình (mean-pooling, mục 2.1–2.2), rồi train tiếp với khoảng 5% compute pretraining ban đầu (abstract; mục 3.1 ghi α = 0,05). Llama 3 dùng GQA với 8 KV head cho cả ba cỡ 8B, 70B và 405B (Bảng 3, mục 3.2 paper Llama 3).

MLA: nén K/V vào một latent

Multi-head latent attention (DeepSeek-V2, 2024) đi theo hướng khác. Thay vì bớt số head, MLA chiếu đầu vào xuống một vector latent nhỏ c_KV và chỉ cache latent đó; K và V của từng head được chiếu lại từ latent khi cần. Phần mã hoá vị trí RoPE không đi qua được phép chiếu này, nên MLA tách ra một phần key nhỏ mang RoPE, cache riêng. Mỗi token mỗi layer chỉ cache d_c + d_h^R giá trị, với d_c = 512 và d_h^R = 64 (mục 3.1.2). Bảng 1 (mục 2.1.4) của paper so sánh KV cache mỗi token:

Chú thích của bảng: KV cache của MLA bằng GQA với chỉ 2,25 nhóm, nhưng chất lượng mạnh hơn MHA. So với DeepSeek 67B, DeepSeek-V2 giảm KV cache 93,3% và tăng throughput sinh tối đa lên 5,76 lần (abstract). Code mẫu 2 ở trên tính đúng công thức này cho DeepSeek-V3: 61 × (512 + 64) × 2 byte ≈ 68,6 KiB mỗi token.

Cái giá: MLA phức tạp hơn, và để chạy nhanh phải “hấp thụ” phép chiếu ngược vào W_Q và W_O, cần kernel riêng. Đổi một model đã train sẵn sang MLA cũng không đơn giản như mean-pooling của GQA.

05PagedAttention và prefix caching

Vấn đề: phân mảnh

Bài 2, mục 8.4 đã nói ý chính của PagedAttention. Đây là con số đứng sau: paper vLLM đo rằng ở các hệ thống trước đó, chỉ 20,4–38,2% bộ nhớ KV cache thực sự chứa trạng thái token (mục 1, Hình 2), nghĩa là khoảng 62–80% bị lãng phí. Ba nguồn lãng phí:

• Dành chỗ trước (reservation): cấp sẵn chỗ cho các token sẽ sinh sau này, chưa dùng tới nhưng không ai khác dùng được.
• Phân mảnh trong (internal fragmentation): cấp theo max_length, request ngắn bỏ trống phần đuôi.
• Phân mảnh ngoài (external fragmentation): các vùng liên tục có kích thước khác nhau để lại khe hở giữa chúng.

Với vLLM, tỉ lệ chứa trạng thái token là 96,3% trong cùng thí nghiệm (Hình 2), và throughput tăng 2–4 lần ở cùng mức latency so với FasterTransformer và Orca (abstract).

Block table

PagedAttention mượn ý tưởng của virtual memory trong hệ điều hành:

• KV cache chia thành block kích thước cố định, mặc định 16 token mỗi block trong vLLM (mục 7.2).
• Mỗi request thấy một dãy block logic liên tục; một block table ánh xạ block logic sang block vật lý nằm bất kỳ đâu trong bộ nhớ GPU.
• Block mới chỉ được cấp khi block cuối đã đầy. Lãng phí tối đa là phần trống của block cuối mỗi sequence.

Kernel attention đọc K, V theo block table thay vì theo một vùng liên tục. Đó là lý do nó có tên riêng: phải viết lại kernel để chịu được việc dữ liệu nằm rải rác.

Chia sẻ block và copy-on-write

Khi block là đơn vị quản lý, việc chia sẻ trở nên rẻ. Mỗi block vật lý có một bộ đếm tham chiếu (ref count). Hai trường hợp quan trọng:

• Parallel sampling: sinh nhiều mẫu cho cùng một prompt (tham số n của API). Các mẫu dùng chung block của prompt. Khi một mẫu cần ghi vào một block đang được chia sẻ, vLLM chép block đó ra một bản mới rồi mới ghi, gọi là copy-on-write ở mức block (mục 4.4, Hình 8).
• Beam search: các beam chia sẻ phần tiền tố chung, và phần chung thay đổi liên tục khi beam bị loại. Bài 9, mục 3 đã nói beam search tốn B lần KV cache; với PagedAttention, phần trùng chỉ lưu một lần.

Paper đo trên OPT-13B với trace Alpaca: chia sẻ block tiết kiệm 6,1–9,8% bộ nhớ cho parallel sampling và 37,6–55,2% cho beam search (mục 6.3, Hình 15).

Prefix caching

Bước tiếp theo là giữ block qua nhiều request. Hai request có cùng system prompt, cùng few-shot example, hay cùng lịch sử hội thoại tới lượt trước, thì K, V của phần prefix đó giống hệt nhau. Không cần prefill lại.

• Automatic prefix caching của vLLM: mỗi block đầy được định danh bằng hash của chính nó cộng hash của block cha, tức là của toàn bộ prefix tính tới block đó. Request mới tra hash; trùng thì tăng ref count và bỏ qua phần prefill tương ứng.
• RadixAttention của SGLang (Zheng et al. 2023, mục 3) giữ một radix tree ánh xạ chuỗi token sang KV cache, với trang nhớ một token, và loại bỏ theo LRU (lá ít dùng gần đây nhất bị loại trước). Kèm scheduling ưu tiên request trùng prefix, SGLang báo throughput cao hơn tới 6,4 lần (abstract).

Prefix caching chỉ giúp TTFT và compute của prefill; nó không làm decode nhanh hơn, và block được giữ lại vẫn chiếm chỗ của KV cache.

Prompt caching phía API

Prompt caching của các API thương mại là cùng một ý, nhìn từ phía hoá đơn: phần prefix đã có K, V trong cache thì nhà cung cấp tính tiền rẻ hơn, vì họ không phải chạy prefill lại. Điều kiện chung là prefix phải trùng khớp từ token đầu tiên: đổi một ký tự trong system prompt là mất cache cho toàn bộ phần sau. Vì vậy nên đặt phần cố định (system prompt, tài liệu, ví dụ) lên đầu và phần thay đổi (câu hỏi của người dùng) xuống cuối. Mục 9 nói về tỉ lệ giá.

06Scheduling: batching, chunked prefill, tách pha

Continuous batching

Bài 2, mục 8.5 đã có ví dụ static batching và continuous batching. Nguồn gốc là Orca (Yu et al., OSDI 2022), với hai ý:

• Iteration-level scheduling: scheduler quyết định ở mức từng bước (iteration) thay vì từng request. Request xong thì rời batch ngay, request mới vào ngay ở bước kế tiếp (mục 3, giải pháp S1).
• Selective batching: các sequence trong batch có độ dài khác nhau, nên không xếp được thành một tensor chữ nhật. Orca chỉ batch những phép không quan tâm độ dài (các phép nhân với weights) và tách riêng attention cho từng sequence (mục 3, giải pháp S2). PagedAttention sau này giải quyết phần attention gọn hơn.

Trên GPT-3 175B, Orca đạt throughput cao hơn FasterTransformer 36,9 lần ở cùng mức latency (abstract; mục 6.2, Hình 10b). Con số lớn như vậy vì baseline phải chờ request dài nhất trong batch xong.

Từ mục 2, ta biết vì sao continuous batching hiệu quả: nó giữ batch luôn đầy, mà batch đầy là cách duy nhất để chia tiền đọc weights cho nhiều token.

Vấn đề còn lại: prefill chen ngang

Continuous batching tạo ra một xung đột mới. Khi một request mới vào, nó cần prefill. Một prefill 8.000 token là hàng trăm mili giây compute; nếu chạy chung bước với các sequence đang decode, mọi người dùng đang stream đều khựng lại trong bước đó. TPOT của họ nhảy vọt. Ngược lại, nếu hoãn prefill để giữ TPOT ổn định, TTFT của request mới tăng.

Chunked prefill

Sarathi-Serve (Agrawal et al. 2024) cắt prefill thành các chunk có lượng compute bằng nhau và chạy dần qua nhiều bước (mục 4.1). Mỗi bước có một ngân sách token: các sequence decode được giữ nguyên trong batch, phần ngân sách còn lại dành cho một chunk prefill, gọi là stall-free batching (mục 4.2). Decode memory-bound còn prefill compute-bound, nên trộn hai thứ trong cùng một bước còn tận dụng được phần compute mà decode bỏ phí. So với vLLM, Sarathi-Serve báo năng lực phục vụ cao hơn 2,6 lần cho Mistral-7B trên một A100, tới 3,7 lần cho Yi-34B trên hai A100 (abstract). Trong vLLM V1, chunked prefill được bật mặc định khi có thể (tài liệu vLLM, mục “Optimization and Tuning”).

Tách prefill và decode ra hai nhóm GPU

Hướng triệt để hơn là disaggregation: prefill chạy trên một nhóm GPU, decode chạy trên nhóm khác, KV cache được chuyển qua mạng sau khi prefill xong. Hai pha không còn tranh nhau, và mỗi nhóm được cấu hình riêng (song song hoá, batch, thậm chí loại GPU). DistServe (Zhong et al. 2024) báo phục vụ được nhiều request hơn 7,4 lần, hoặc SLO chặt hơn 12,6 lần, mà vẫn giữ hơn 90% request trong giới hạn latency (abstract). Splitwise (Patel et al. 2023) báo throughput cao hơn 1,4 lần với chi phí thấp hơn 20%, hoặc cao hơn 2,35 lần với cùng ngân sách điện và chi phí (abstract). Cái giá là phải chuyển KV cache giữa các máy, cần mạng nhanh và hệ thống phức tạp hơn nhiều.

07FlashAttention: giảm byte, không giảm FLOPs

Bài 2, mục 8.6 và bài 7, mục 10 đã nói FlashAttention tính theo block trong SRAM, không dựng ma trận attention đầy đủ trong HBM, và cho kết quả chính xác chứ không xấp xỉ. Phần chưa nói là làm sao tính softmax theo block được, khi softmax cần tổng của cả hàng.

Online softmax

Softmax ổn định số học cần hai thứ của cả hàng: giá trị lớn nhất m (để trừ đi, tránh tràn số) và tổng l của exp(x − m). Cách ngây thơ đọc hàng ba lần: một lần tìm max, một lần tính tổng, một lần chuẩn hoá. Online softmax (Milakov & Gimelshein 2018, Thuật toán 3 ở mục 3) gộp hai lần đầu làm một bằng cách giữ max và tổng chạy, và sửa lại tổng mỗi khi max thay đổi:

m_j = max(m_{j-1}, x_j)
l_j = l_{j-1} · exp(m_{j-1} − m_j) + exp(x_j − m_j)

Thừa số exp(m_{j-1} − m_j) co phần tổng cũ về thang của max mới. Paper này đo softmax nhanh hơn tới 1,3 lần chỉ nhờ bớt một lần đọc bộ nhớ (abstract).

FlashAttention áp cùng ý cho cả output. Duyệt K, V theo từng block; với mỗi block, tính score của block đó, cập nhật m và l, rồi co output đang tích luỹ bằng cùng thừa số trước khi cộng phần mới:

với mỗi block j của K, V:
  S   = Q · K_jᵀ / √d                      (chỉ n × block, nằm trong SRAM)
  m'  = max(m, rowmax(S))
  P   = exp(S − m')
  l   = l · exp(m − m') + rowsum(P)
  O   = O · exp(m − m') + P · V_j
  m   = m'
cuối cùng: O = O / l

Code mẫu 3a: online softmax so với attention chuẩn

Viết lại thuật toán trên bằng PyTorch thuần, chạy trên CPU, rồi so với attention chuẩn (dựng đủ ma trận n × n) và với scaled_dot_product_attention của PyTorch. Bản fp64 dùng làm tham chiếu.

# Python 3.12, torch 2.14.1 — chạy trên CPU, fp32
import torch, torch.nn.functional as F
torch.manual_seed(0)

def attention_standard(q, k, v):
    s = q @ k.transpose(-1, -2) / q.shape[-1] ** 0.5       # dựng đủ ma trận n × n
    mask = torch.ones(s.shape[-2:], dtype=torch.bool).triu(1)
    return torch.softmax(s.masked_fill(mask, float("-inf")), dim=-1) @ v

def attention_online(q, k, v, block=128):
    """Duyệt K, V theo block; mỗi hàng query chỉ giữ max m, tổng l và output o chạy."""
    n, d = q.shape[-2], q.shape[-1]
    o = torch.zeros_like(q)
    m = torch.full(q.shape[:-1] + (1,), float("-inf"))
    l = torch.zeros(q.shape[:-1] + (1,))
    rows = torch.arange(n).unsqueeze(1)
    for j in range(0, n, block):
        kj, vj = k[..., j:j + block, :], v[..., j:j + block, :]
        s = q @ kj.transpose(-1, -2) / d ** 0.5            # chỉ n × block
        s = s.masked_fill(torch.arange(j, j + kj.shape[-2]) > rows, float("-inf"))
        m_new = torch.maximum(m, s.amax(-1, keepdim=True))
        p = torch.exp(s - m_new)
        scale = torch.exp(m - m_new)                        # chỉnh lại phần đã cộng theo max mới
        l = l * scale + p.sum(-1, keepdim=True)
        o = o * scale + p @ vj
        m = m_new
    return o / l

q, k, v = (torch.randn(8, 2048, 128) for _ in range(3))   # 8 head, 2.048 token, head_dim 128
ref64 = attention_standard(q.double(), k.double(), v.double())
std = attention_standard(q, k, v)
onl = attention_online(q, k, v)
sdpa = F.scaled_dot_product_attention(q, k, v, is_causal=True)
for name, x in [("standard fp32", std), ("online fp32", onl), ("SDPA fp32", sdpa)]:
    print(f"{name:14} max |x - ref_fp64| = {(x.double() - ref64).abs().max():.2e}")
print(f"online vs standard: max diff {(onl - std).abs().max():.2e}, "
      f"allclose(atol=1e-5) = {torch.allclose(onl, std, atol=1e-5)}")
n, b = 2048, 128
print(f"phần tử score giữ cùng lúc mỗi head: standard {n*n:,} · online {n*b:,} ({n*n // (n*b)}× ít hơn)")

Output:

standard fp32  max |x - ref_fp64| = 8.70e-07
online fp32    max |x - ref_fp64| = 9.05e-07
SDPA fp32      max |x - ref_fp64| = 1.00e-06
online vs standard: max diff 5.96e-07, allclose(atol=1e-5) = True
phần tử score giữ cùng lúc mỗi head: standard 4,194,304 · online 262,144 (16× ít hơn)

Ba cách tính đều lệch khỏi bản fp64 cỡ 1e-6, đúng mức sai số làm tròn của fp32. Online softmax không phải xấp xỉ: nó là cùng một phép tính với thứ tự cộng khác, và như bài 9, mục 7 (tầng 2) đã nói, thứ tự cộng khác thì kết quả lệch ở chữ số cuối. Dòng cuối cho thấy cái lợi: mỗi lúc chỉ cần giữ một khối n × 128 thay vì cả ma trận n × n. Trên GPU, khối đó nằm vừa trong SRAM, nên ma trận attention không bao giờ phải ghi ra HBM.

Độ phức tạp IO

Định lý 2 (mục 3.2) của paper FlashAttention đếm số lần truy cập HBM, với N là độ dài chuỗi, d là head dim và M là kích thước SRAM:

attention chuẩn:  Θ(N·d + N²)   lần truy cập HBM
FlashAttention:   Θ(N²·d²/M)    lần truy cập HBM

Với d cỡ 64–128 và M cỡ 100 KB, d²/M nhỏ hơn 1 nhiều lần, nên số byte đọc ghi giảm mạnh. FLOPs thì không giảm: vẫn là đủ QKᵀ và PV. Bộ nhớ phụ chỉ tăng tuyến tính theo N (Định lý 1, mục 3.1). Ở backward, FlashAttention còn tính nhiều FLOPs hơn: nó chỉ lưu output và hai thống kê m, l, rồi tính lại ma trận attention từ Q, K, V khi cần (mục 3.1, đoạn “Recomputation”). Đổi thêm compute lấy ít byte hơn là có lời, vì attention memory-bound. Kết quả: attention trên GPT-2 nhanh hơn 7,6 lần so với bản PyTorch (Hình 1), train GPT-2 với context 1K nhanh hơn 3 lần so với bản HuggingFace và BERT-large nhanh hơn 15% so với kỷ lục MLPerf 1.1 (abstract).

FlashAttention-2, FlashAttention-3 và FlashDecoding

• FlashAttention-2 (Dao 2023) chia việc giữa các thread block và warp tốt hơn, bớt phép tính không phải nhân ma trận; khoảng 2 lần nhanh hơn bản đầu, đạt 50–73% FLOPs lý thuyết của A100 (abstract).
• FlashAttention-3 (Shah et al. 2024) dành cho H100, tận dụng xử lý bất đồng bộ của Hopper và FP8: nhanh hơn FlashAttention-2 1,5–2,0 lần ở forward, tới 740 TFLOPs/s với FP16 (75% công suất) và gần 1,2 PFLOPs/s với FP8 (abstract, mục 1).
• Flash-Decoding (Dao, Haziza, Massa, Sizov 2023) giải quyết riêng pha decode. Ở decode chỉ có một query mỗi sequence, mà FlashAttention song song hoá theo query và batch, nên với batch 1 nó dùng chưa tới 1% GPU. Flash-Decoding chia K, V thành nhiều đoạn, tính attention của query với từng đoạn song song, lưu thêm log-sum-exp của mỗi đoạn, rồi gộp lại bằng chính phép co của online softmax. Bài viết báo sinh nhanh hơn tới 8 lần với chuỗi rất dài.

Lưu ý phần cứng: Turing không có FlashAttention-2

README của flash-attn ghi rõ: FlashAttention-2 cần GPU Ampere, Ada hoặc Hopper (A100, RTX 3090, RTX 4090, H100); với Turing (T4, RTX 2080) phải dùng repo riêng flash-attention-turing, chỉ hỗ trợ một phần tính năng. Trong PyTorch, backend FLASH_ATTENTION của scaled_dot_product_attention chỉ bật cho kiến trúc từ sm80 trở lên, còn backend memory-efficient (dựa trên xFormers) chạy từ sm50 (theo sdp_utils.cpp trong source PyTorch). RTX 2080 Ti là sm75, nên khi gọi SDPA, PyTorch tự chọn memory-efficient, không phải flash. Kernel memory-efficient cũng tính theo block và không dựng ma trận n × n, nên lợi ích về bộ nhớ vẫn có.

Code mẫu 3b: so backend math và memory-efficient

# Code minh hoạ. Python 3.12, torch 2.14.1+cu130
import time, torch, torch.nn.functional as F
from torch.nn.attention import sdpa_kernel, SDPBackend

def bench(backend, n, heads=12, d=128, iters=10):
    q, k, v = (torch.randn(1, heads, n, d, device="cuda", dtype=torch.float16) for _ in range(3))
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    base = torch.cuda.memory_allocated()
    with sdpa_kernel(backend):                      # ép PyTorch chỉ dùng backend này
        F.scaled_dot_product_attention(q, k, v, is_causal=True)   # warmup
        torch.cuda.synchronize(); t0 = time.perf_counter()
        for _ in range(iters):
            F.scaled_dot_product_attention(q, k, v, is_causal=True)
        torch.cuda.synchronize()
    dt = (time.perf_counter() - t0) / iters
    return dt, (torch.cuda.max_memory_allocated() - base) / 2**20

for n in (1024, 2048, 4096, 8192, 16384):
    for backend in (SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION):
        try:
            dt, peak = bench(backend, n)
            print(f"n={n:6d} {backend.name:20s} {dt * 1e3:8.2f} ms   bộ nhớ đỉnh thêm {peak:8.0f} MiB")
        except torch.OutOfMemoryError:
            print(f"n={n:6d} {backend.name:20s} hết bộ nhớ")
        torch.cuda.empty_cache()

Chưa chạy thử. Cần GPU NVIDIA có CUDA (ví dụ RTX 2080 Ti 11 GB, Turing sm_75), CUDA 13.0, Python 3.12, torch 2.14.1+cu130. Trên Turing, thêm SDPBackend.FLASH_ATTENTION vào vòng lặp sẽ báo lỗi không có kernel phù hợp, vì backend này cần sm80 trở lên.

Điều nên thấy khi chạy: bộ nhớ đỉnh của MATH tăng theo n² vì nó dựng ma trận score heads × n × n (ở n = 16.384 và 12 head, riêng ma trận fp16 đó đã khoảng 6 GiB), còn của EFFICIENT_ATTENTION chỉ tăng tuyến tính.

08Speculative decoding

Ý tưởng

Mục 2 cho thấy decode lãng phí compute: đọc toàn bộ weights chỉ để tính một token. Batching tận dụng phần thừa bằng cách phục vụ nhiều người dùng. Speculative decoding tận dụng nó cho một người dùng: đoán trước vài token, rồi kiểm cả loạt trong một lượt forward.

Một vòng gồm ba bước:

1. Model nháp (draft model), nhỏ và nhanh, sinh tuần tự k token đoán.
2. Model đích (target model) chạy một lượt forward trên prefix cộng k token đoán, ra phân phối ở cả k + 1 vị trí. Lượt này giống một prefill ngắn: đọc weights một lần cho k + 1 vị trí, nên khi decode memory-bound thì nó tốn gần bằng một bước decode thường.
3. Bước kiểm quyết định giữ bao nhiêu token đoán, rồi luôn thêm ít nhất một token từ model đích.

Vì sao output giữ đúng phân phối

Gọi p(x) là phân phối của model đích và q(x) của model nháp tại một vị trí. Với token nháp x ~ q, quy tắc của Leviathan et al. (mục 2.3) là:

• Nếu q(x) ≤ p(x): luôn nhận.
• Nếu q(x) > p(x): nhận với xác suất p(x)/q(x), tức nhận với xác suất min(1, p(x)/q(x)) trong mọi trường hợp.
• Nếu bị từ chối: lấy mẫu lại từ p′(x) = norm(max(0, p(x) − q(x))), rồi dừng vòng.

Chứng minh ngắn (Phụ lục A.1 của paper): xác suất ra đúng token x′ là tổng của hai đường.

P(nhận, x = x′)       = q(x′) · min(1, p(x′)/q(x′)) = min(p(x′), q(x′))
P(từ chối, rồi ra x′) = (1 − Σ_x min(p(x), q(x))) · p′(x′) = p(x′) − min(p(x′), q(x′))
Cộng lại:             P(x = x′) = p(x′)

Dòng thứ hai đúng vì Σ_x max(0, p − q) = 1 − Σ_x min(p, q), nên hằng số chuẩn hoá của p′ triệt tiêu với xác suất từ chối. Kết quả: token ra có phân phối đúng bằng p, không phải xấp xỉ. Chen et al. 2023 (DeepMind) chứng minh cùng kết quả độc lập (Định lý 1 ở phụ lục), với ký hiệu ngược lại: ở paper đó q là model đích còn p là model nháp. Với greedy decoding, quy tắc rút gọn thành: nhận token nháp nếu nó trùng argmax của model đích.

Có một điểm tinh tế, nối với bài 9, mục 7 (tầng 2): “giống hệt về phân phối” khác với “giống từng bit”. Model đích tính k + 1 vị trí trong một lượt thay vì từng vị trí một, nên kernel chia phép tính khác đi, và logits có thể lệch ở chữ số cuối. Với greedy, khi hai token đứng đầu gần bằng nhau, output speculative có thể khác output thường ở vài token. Đó là chuyện số thực, không phải lỗi của thuật toán.

Mỗi vòng được bao nhiêu token

Leviathan et al. định nghĩa tỉ lệ chấp nhận β là xác suất nhận một token nháp, và α = E(β) (mục 3.1). Giả sử các lần nhận độc lập, số token kỳ vọng mỗi vòng là (Công thức 1, mục 3.1, paper viết γ thay cho k):

E[số token mỗi vòng] = (1 − α^(k+1)) / (1 − α)

Ví dụ với k = 4: α = 0,6 cho khoảng 2,3 token mỗi vòng; α = 0,8 cho khoảng 3,4; α = 0,9 cho khoảng 4,1. Tốc độ thật còn phải trừ thời gian chạy model nháp, nên Định lý 3.8 (mục 3.3) đưa thêm tỉ lệ chi phí c giữa model nháp và model đích. Model nháp càng rẻ và càng giống model đích thì càng lời; k lớn chỉ có lợi khi α cao.

Trên T5-XXL, paper báo nhanh hơn 2–3 lần với output giống hệt (abstract; Bảng 2 có các mức từ 1,4 đến 3,4 lần). Chen et al. báo nhanh hơn 2–2,5 lần khi lấy mẫu từ Chinchilla 70B (abstract).

Medusa và EAGLE: bỏ model nháp riêng

Chạy hai model có hai cái khó: phải có một model nhỏ cùng tokenizer, và phải quản lý hai KV cache.

• Medusa (Cai et al. 2024) gắn thêm vài decoding head lên chính model đích, mỗi head đoán token ở vị trí xa hơn, rồi kiểm nhiều nhánh ứng viên cùng lúc bằng tree attention. Medusa-1 train head trên backbone đóng băng, nhanh hơn trên 2,2 lần mà không giảm chất lượng; Medusa-2 train cả backbone, nhanh hơn 2,3–2,8 lần (PDF v3, abstract và mục 1); con số 3,6 lần chỉ là category tốt nhất của MT-Bench. Lưu ý: Medusa-2 đổi weights của backbone, và cơ chế “typical acceptance” chỉ giữ chất lượng tương tự chứ không giữ đúng phân phối như rejection sampling.
• EAGLE (Li et al. 2024) đoán ở mức feature (đầu ra của layer gần cuối) thay vì ở mức token, bằng một head nhẹ, và vẫn kiểm bằng rejection sampling nên giữ đúng phân phối. Với LLaMA2-Chat 70B, EAGLE báo latency giảm 2,7–3,5 lần và throughput gấp đôi (abstract; mục 4.4, Bảng 7).

Code mẫu 4: assisted generation của Transformers

Hugging Face Transformers gọi speculative decoding là assisted generation: truyền một model nhỏ vào tham số assistant_model của generate(). Tài liệu ghi rõ nó hỗ trợ greedy và sampling nhưng không hỗ trợ input theo batch. Thư viện không trả về tỉ lệ chấp nhận, nên code dưới đếm số lượt forward của model đích bằng một forward hook: số token sinh ra chia cho số vòng kiểm là số token trung bình mỗi vòng, từ đó suy ra α theo công thức ở trên.

# Code minh hoạ. Python 3.12, torch 2.14.1+cu130, transformers 5.18.0
import time, torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct")
load = dict(dtype=torch.float16, device_map="cuda")
target = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct", **load).eval()
draft = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", **load).eval()  # cùng tokenizer

calls = 0
def count_call(module, args, output):
    global calls
    calls += 1
target.register_forward_hook(count_call)            # mỗi lượt forward của model đích = một vòng kiểm

prompts = {
    "code": "Viết hàm Python kiểm tra một số có phải số nguyên tố không, kèm docstring.",
    "văn xuôi": "Kể một câu chuyện ngắn về một người bán phở ở Hà Nội vào buổi sáng mùa đông.",
}
K = 5
for name, text in prompts.items():
    inputs = tok.apply_chat_template([{"role": "user", "content": text}], add_generation_prompt=True,
                                     return_tensors="pt", return_dict=True).to("cuda")
    gen = dict(max_new_tokens=256, do_sample=False)

    torch.cuda.synchronize(); t0 = time.perf_counter()
    plain = target.generate(**inputs, **gen)
    torch.cuda.synchronize(); t_plain = time.perf_counter() - t0

    calls = 0
    torch.cuda.synchronize(); t0 = time.perf_counter()
    spec = target.generate(**inputs, **gen, assistant_model=draft,
                           num_assistant_tokens=K, num_assistant_tokens_schedule="constant")
    torch.cuda.synchronize(); t_spec = time.perf_counter() - t0

    n_new = spec.shape[1] - inputs["input_ids"].shape[1]
    per_round = n_new / calls                       # vòng đầu xử lý luôn cả prompt, không có lượt prefill riêng
    print(f"{name:8s}: trùng output thường = {torch.equal(plain, spec)}, "
          f"nhanh hơn {t_plain / t_spec:.2f}×, {per_round:.2f} token mỗi vòng kiểm")

Chưa chạy thử. Cần GPU NVIDIA còn trống khoảng 6 GB VRAM trở lên (ví dụ RTX 2080 Ti 11 GB), CUDA 13.0, Python 3.12, torch 2.14.1+cu130, transformers 5.18.0. Model nháp và model đích phải dùng chung tokenizer (Qwen2.5-0.5B và 1.5B dùng chung).

Điều nên thấy khi chạy: prompt code thường cho nhiều token mỗi vòng hơn prompt văn xuôi tiếng Việt, vì code có nhiều đoạn dễ đoán (thụt lề, từ khoá, tên biến lặp lại) mà model nhỏ đoán trúng. Cột “trùng output thường” có thể là False ở một số prompt vì lý do số thực đã nói ở trên, không phải vì thuật toán sai. Ngoài ra, transformers có thể dừng nháp sớm khi model nháp không đủ tự tin (assistant_confidence_threshold), nên số token nháp mỗi vòng có thể ít hơn K.

09Context dài đắt ở đâu

Câu trả lời quen thuộc là “attention O(n²)”. Câu đó đúng một phần. Tách chi phí của context dài ra ba khoản thì thấy rõ hơn.

Khoản 1: prefill, nơi O(n²) thật sự có mặt

Dùng phép đếm FLOPs của bài 8, mục 6: mỗi token tốn 2N cho các ma trận cộng 2 · n_layer · n_ctx · d_attn cho phần attention phụ thuộc context. Với prefill n token có causal mask, mỗi token trung bình nhìn n/2 token phía trước, và cả QKᵀ lẫn nhân với V đều tốn 2 FLOP mỗi cặp mỗi chiều, nên phần attention trung bình mỗi token đúng bằng số hạng của bài 8 với n_ctx = n. Attention bắt đầu lớn hơn phần ma trận khi:

2 · n_layer · n · d_attn > 2N   ⇔   n > N / (n_layer · d_attn)

Với Llama 3.1 8B (N không tính embedding 6,98B, 32 layer, d_attn 4.096), điểm cắt là khoảng 53.000 token, đúng như dòng cuối output của code mẫu 2. Ở context 4k, attention chỉ chiếm 7% FLOPs của prefill; ở 32k là 38%; ở 128k là 71%. Vậy O(n²) là khoản đáng kể chỉ khi context rất dài, và nó nằm ở prefill, tức là ở TTFT.

Khoản 2: decode, nơi KV cache là thủ phạm

Mỗi bước decode đọc toàn bộ KV cache của sequence. Với Llama 3.1 8B, đó là 0,54 GB ở context 4k và 17 GB ở context 128k, so với 16 GB weights. Nghĩa là:

• TPOT tăng tuyến tính theo context: ở 128k, mỗi bước decode đọc gấp đôi lượng byte so với context ngắn, chỉ riêng một sequence.
• Batch tối đa giảm theo context: trên GPU 80 GB, từ 104 sequence ở 4k xuống 3 sequence ở 128k (code mẫu 2). Throughput giảm theo.

Đây là khoản phải trả ở mọi token sinh ra, và nó tuyến tính chứ không bậc hai. Với phần lớn workload thực tế (prompt vài nghìn tới vài chục nghìn token, output dài), khoản này lớn hơn khoản 1. Đó là lý do GQA, MLA, KV cache quantization và PagedAttention, những kỹ thuật cắt byte, quan trọng với context dài không kém FlashAttention.

Bảng tổng hợp cho Llama 3.1 8B, bf16, trên một GPU 80 GB (số lấy từ output code mẫu 2):

Khoản 3: chất lượng

Context window công bố là giới hạn kỹ thuật, không phải độ dài model dùng tốt. RULER (Hsieh et al. 2024) đánh giá 17 model long-context trên 13 tác vụ: tất cả đều công bố context từ 32K token trở lên, nhưng chỉ một nửa giữ được kết quả chấp nhận được ở độ dài 32K (abstract; độ dài hiệu quả ở Bảng 3). Context dài vừa tốn tiền vừa có thể không mang lại thông tin model thực sự dùng được.

Một hướng hay bị hiểu nhầm là attention sink và sliding window. StreamingLLM (Xiao et al. 2023) quan sát rằng model dồn rất nhiều attention vào vài token đầu tiên, dù chúng không mang nghĩa gì, và gọi chúng là “attention sink”. Giữ K, V của 4 token đầu cộng một cửa sổ trượt các token gần nhất (mục 1, mục 3.2) cho phép model sinh ổn định tới 4 triệu token và nhanh hơn tới 22,2 lần so với sliding window tính lại (abstract). Nhưng paper nói rõ phương pháp này không mở rộng context window và không tăng trí nhớ dài hạn (phụ lục A): token đã rời cửa sổ là bị quên. Nó là cách giữ bộ nhớ cố định, không phải cách nhớ được nhiều hơn.

Giá API phản ánh ba khoản trên

Bảng giá của các nhà cung cấp lớn (lấy ngày 08/10/2026 từ trang giá của OpenAI, Anthropic và Google Gemini) có cùng một hình dạng. Bài chỉ nêu tỉ lệ, vì giá tuyệt đối đổi nhanh:

• Output đắt hơn input khoảng 4–8 lần, phổ biến nhất là 5–6 lần ở các model flagship hiện tại. Lý do nằm ở mục 2: một token input đi qua prefill compute-bound, chia chung lần đọc weights với hàng nghìn token khác; một token output là một bước decode memory-bound, chiếm chỗ KV cache suốt thời gian sinh.
• Cached input rẻ hơn input thường, thường còn khoảng 10%. Có ngoại lệ: các model mới nhất của Anthropic còn 2,5–5%, vài model cũ của OpenAI là 25–50%. Đây là prefix caching ở mục 5: phần prefix đã có K, V thì không phải prefill lại. Anthropic tính thêm phí ghi cache (1,25 lần giá input cho cache 5 phút, 2 lần cho cache 1 giờ), còn Gemini tính thêm phí lưu trữ cache theo giờ, vì block được giữ lại chiếm bộ nhớ GPU.
• Prompt rất dài có thể có giá riêng. Một số model Gemini Pro tính đắt hơn cho prompt trên 200k token, OpenAI có mức giá long context cho input trên 272K token. Đây là khoản 1 và khoản 2 cộng lại: prefill đắt hơn và mỗi request chiếm phần lớn bộ nhớ của GPU.

10Những chỗ dễ hiểu sai

Chín hiểu lầm dưới đây khớp với hàng trên và hàng dưới của Hình 1.

Tổng kết: mỗi byte đọc vào làm được bao nhiêu việc

Toàn bộ bài này quy về một tỉ số: số FLOP làm được cho mỗi byte đọc từ bộ nhớ.

• Prefill xử lý cả prompt trong một lượt, mỗi byte weights phục vụ hàng nghìn token, nên compute-bound và quyết định TTFT. Decode sinh từng token, mỗi bước đọc lại toàn bộ weights và KV cache để làm khoảng 1 FLOP mỗi byte ở batch 1, trong khi RTX 2080 Ti cần khoảng 87 FLOP mỗi byte mới dùng hết compute.
• Batching chia lần đọc weights cho nhiều sequence, nên throughput tăng gần tuyến tính theo batch. Nhưng KV cache không chia được: mỗi sequence có K, V riêng. Vì vậy bộ nhớ dành cho KV cache quyết định batch tối đa, và batch quyết định throughput.
• KV cache tính được từ config.json: Llama 3.1 8B tốn 128 KiB mỗi token, một sequence 128k nặng ngang cả bản weights. GQA, MQA và MLA giảm số byte đó bằng cách chia sẻ hoặc nén K, V; KV cache quantization giảm byte mỗi giá trị.
• PagedAttention đưa tỉ lệ bộ nhớ KV thực sự dùng từ 20–38% lên 96%, và biến việc chia sẻ prefix thành chuyện tăng ref count. Prefix caching và prompt caching của API là cùng một ý.
• Continuous batching giữ batch luôn đầy; chunked prefill và tách prefill/decode giải quyết xung đột giữa TTFT của request mới và TPOT của request đang chạy.
• FlashAttention giữ nguyên FLOPs, giảm byte nhờ online softmax và tiling trong SRAM; kết quả chính xác, chỉ khác thứ tự cộng. Trên Turing, PyTorch dùng kernel memory-efficient thay cho flash.
• Speculative decoding dùng compute thừa của decode để kiểm nhiều token đoán trong một lượt; rejection sampling giữ đúng phân phối của model đích.
• Context dài đắt chủ yếu ở bộ nhớ và băng thông của KV cache, phải trả ở mọi token output; O(n²) chỉ trội ở prefill khi context rất dài. Giá API phản ánh đúng điều đó: output đắt hơn input 4–8 lần, cached input thường chỉ còn khoảng 10%.

Bài 14 sẽ nói về Mixture of Experts: tổng weights rất lớn nhưng mỗi token chỉ dùng một phần, nên tỉ lệ giữa bộ nhớ và compute đổi hẳn so với model dense, và mọi phép tính roofline trong bài này phải làm lại. Còn reasoning model sinh output rất dài, nên decode chiếm phần lớn chi phí của chúng; bài 15 sẽ quay lại chuyện đó.

Tài liệu tham khảo

Paper

• Williams, Waterman & Patterson 2009 — Roofline: An Insightful Visual Performance Model for Multicore Architectures
• Milakov & Gimelshein 2018 — Online normalizer calculation for softmax
• Shazeer 2019 — Fast Transformer Decoding: One Write-Head is All You Need (MQA)
• Yu et al. 2022 — Orca: A Distributed Serving System for Transformer-Based Generative Models (OSDI 2022)
• Dao et al. 2022 — FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Pope et al. 2022 — Efficiently Scaling Transformer Inference
• Leviathan, Kalman & Matias 2022 — Fast Inference from Transformers via Speculative Decoding
• Chen et al. 2023 — Accelerating Large Language Model Decoding with Speculative Sampling
• Ainslie et al. 2023 — GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
• Dao 2023 — FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• Kwon et al. 2023 — Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM)
• Xiao et al. 2023 — Efficient Streaming Language Models with Attention Sinks (StreamingLLM)
• Patel et al. 2023 — Splitwise: Efficient Generative LLM Inference Using Phase Splitting
• Zheng et al. 2023 — SGLang: Efficient Execution of Structured Language Model Programs (RadixAttention)
• Cai et al. 2024 — Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads
• Zhong et al. 2024 — DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving
• Li et al. 2024 — EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty
• Liu et al. 2024 — KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache
• Agrawal et al. 2024 — Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve
• Hsieh et al. 2024 — RULER: What’s the Real Context Size of Your Long-Context Language Models? (bản v3)
• DeepSeek-AI 2024 — DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (MLA)
• Llama Team 2024 — The Llama 3 Herd of Models
• Shah et al. 2024 — FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
• DeepSeek-AI 2024 — DeepSeek-V3 Technical Report

Blog, tài liệu và trang giá

• NVIDIA — NVIDIA Turing GPU Architecture Whitepaper (Bảng 1)
• kipply 2022 — Transformer Inference Arithmetic
• Dao, Haziza, Massa & Sizov 2023 — Flash-Decoding for long-context inference
• Dao-AILab — flash-attention README
• PyTorch — torch.nn.functional.scaled_dot_product_attention, torch.nn.attention.SDPBackend, sdp_utils.cpp
• vLLM — Automatic Prefix Caching, Optimization and Tuning (chunked prefill), Quantized KV Cache
• Hugging Face — Cache strategies, Speculative decoding (assisted decoding), Assisted Generation
• Trang giá (lấy ngày 08/10/2026) — OpenAI API pricing, Anthropic Claude pricing, Gemini API pricing