Suy luận & Tối ưu hóa (Inference): Giải mã nhanh, bộ nhớ đệm KV và giới hạn phần cứng

Khi một mô hình ngôn ngữ đã được huấn luyện xong và đóng băng toàn bộ trọng số $\theta$, công việc duy nhất của nó tại mỗi bước là nhận vào chuỗi ngữ cảnh $x_1, x_2, \dots, x_t$ và xuất ra một vector điểm số chưa chuẩn hóa $\mathbf{z} \in \mathbb{R}^{|V|}$ (gọi là logits) trên toàn bộ từ vựng $V$.

Tưởng chừng mọi thứ đến đây đã xong, nhưng để biến vector logits đó thành một hệ thống sinh văn bản vừa thông minh, tự nhiên, vừa đạt tốc độ phản hồi tức thì, chúng ta phải giải quyết hai câu hỏi lớn định hình toàn bộ kỹ thuật suy luận LLM hiện đại:

• Làm thế nào để chọn ra token tiếp theo từ phân phối xác suất $p(x_{t+1} \mid x_{1:t})$ sao cho văn bản sinh ra không bị rơi vào vòng lặp vô hồn của tìm kiếm tham lam (greedy decoding), cũng không bị trượt sang những từ vô nghĩa ở phần đuôi dài của phân phối?
• Tại sao khi sinh từng token tự hồi quy, các siêu chip GPU hàng trăm TFLOPS lại dành tới hơn 95% thời gian chỉ để… đứng chờ bộ nhớ RAM, và làm thế nào để phá vỡ bức tường băng thông bộ nhớ (memory wall) này mà không làm giảm chất lượng mô hình?

1. Chiến lược giải mã và chọn mẫu (Decoding & Sampling)

Giả sử mô hình vừa đọc xong cụm từ “Cánh cửa gỗ bất ngờ mở ra, bên trong là một” và trả về vector logits $\mathbf{z} = (z_1, z_2, \dots, z_{|V|})$ ở lớp tuyến tính cuối cùng. Hàm Softmax biến vector này thành phân phối xác suất của token tiếp theo:

Có phân phối xác suất rồi, ta rút ra từ tiếp theo bằng cách nào? Quyết định tưởng như nhỏ nhặt này lại có ảnh hưởng sống còn đến chất lượng câu trả lời của mô hình.

1.1 Giải mã cực đại xác suất và căn bệnh thoái hóa văn bản

Nếu mục tiêu là tìm chuỗi từ $y_1, \dots, y_M$ có xác suất hợp lý cao nhất theo mô hình, phản xạ tự nhiên nhất là giải mã tham lam (Greedy Decoding): tại mỗi bước $t$, ta luôn chọn từ có xác suất lớn nhất:

Giải mã tham lam hoàn toàn tất định (deterministic) và rất phù hợp cho các tác vụ có đáp án duy nhất (như phân loại hay trích xuất số liệu). Tuy nhiên, việc chọn cực đại cục bộ ở từng bước không đảm bảo chuỗi tổng thể có xác suất cao nhất: một từ hơi kém phổ biến hơn ở bước $t$ có thể mở ra một nhánh câu xuất sắc ở bước $t+1$.

Để khắc phục, các hệ thống dịch máy cổ điển (Chương 0, mục 9) sử dụng tìm kiếm chùm (Beam Search): duy trì song song $B$ ứng viên chuỗi tốt nhất (gọi là độ rộng chùm beam width) ở mỗi bước và chọn chuỗi tối đa hóa $\prod_{t} p(x_t \mid x_{<t})$.

Thế nhưng, khi áp dụng Greedy Decoding hoặc Beam Search vào sinh văn bản mở (open-ended generation) trên các mô hình ngôn ngữ lớn, một hiện tượng kỳ lạ xảy ra: chỉ sau vài câu, mô hình bắt đầu viết nhạt nhẽo, an toàn quá mức, và nhanh chóng rơi vào những vòng lặp vô tận kiểu: “Tôi không biết. Tôi không biết. Tôi không biết…”

Năm 2020, Ari Holtzman và cộng sự đã đặt tên cho hiện tượng này là sự thoái hóa văn bản nơ-ron (neural text degeneration) (Holtzman & others, 2020). Vậy nếu không lấy cực đại (argmax), tại sao ta không lấy mẫu ngẫu nhiên trực tiếp từ phân phối $p(w \mid x_{1:t})$ như cách chúng ta từng làm với mô hình xích Markov ở Chương 0?

Câu trả lời nằm ở cái đuôi không đáng tin cậy (unreliable tail) của hàm Softmax. Với bộ từ vựng $|V| \approx 128{,}000$ token, luôn có hàng chục nghìn từ hoàn toàn lạc đề hoặc sai ngữ pháp nhưng vẫn nhận một xác suất nhỏ khác không (ví dụ $10^{-5}$). Khi cộng dồn $100{,}000$ từ như vậy lại, cái đuôi dài này có thể chiếm tới $10\% - 20\%$ tổng khối lượng xác suất! Nếu lấy mẫu thuần túy trên toàn bộ từ vựng, cứ sau khoảng 5 đến 10 từ, mô hình sẽ trượt chân rút phải một từ hoàn toàn vô nghĩa ở phần đuôi, kéo toàn bộ các bước tự hồi quy phía sau chệch hướng.

1.2 Nhiệt độ ($T$), Top-$k$ và Chọn mẫu Hạt nhân Top-$p$ (Nucleus Sampling)

Để vừa giữ được sự ngẫu nhiên sáng tạo, vừa cắt bỏ hoàn toàn cái đuôi nhiễu loạn, các hệ thống LLM hiện đại kết hợp một đường ống gồm ba phép biến đổi liên tiếp trên vector logits $\mathbf{z}$ trước khi rút mẫu:

1. Nhiệt độ (Temperature Scaling — $T > 0$): Trước khi đưa vào Softmax, ta chia toàn bộ logits cho hằng số nhiệt độ $T$:

   $$p_i(T) = \frac{\exp(z_i / T)}{\sum_{j=1}^{|V|} \exp(z_j / T)}.$$

   Khi $T < 1$, chênh lệch giữa các logit lớn và nhỏ được khuếch đại, làm phân phối nhọn hơn (tập trung vào các từ an toàn; khi $T \to 0$, nó hội tụ về giải mã tham lam). Khi $T > 1$, phân phối bị san phẳng, trao nhiều cơ hội hơn cho các từ ít phổ biến (nhưng cũng làm phình to cái đuôi nhiễu).

2. Lọc Top-$k$ (Top-$k$ Sampling): Sắp xếp các token theo thứ tự xác suất giảm dần $p_{(1)} \ge p_{(2)} \ge \dots \ge p_{(|V|)}$, chỉ giữ lại đúng $k$ token đầu tiên (ví dụ $k = 50$) và gán logit của mọi token còn lại bằng $-\infty$. Nhược điểm của Top-$k$ là số lượng $k$ cố định không thích ứng được với ngữ cảnh:

   • Khi phân phối rất nhọn (ví dụ: “Thủ đô của Pháp là…”, từ "Paris" chiếm $98\%$), việc giữ $k=50$ từ sẽ vô tình đưa 49 từ sai vào tập chọn mẫu.
   • Khi phân phối rất phẳng (ví dụ: “Hôm nay tôi muốn kể câu chuyện về một…”, có hàng trăm danh từ hợp lý mỗi từ chiếm $0.5\%$), việc cắt ở $k=10$ hay $k=20$ lại chặt cụt mất phần lớn các ý tưởng hay.

3. Chọn mẫu Hạt nhân Top-$p$ (Nucleus Sampling [@holtzman2020]): Thay vì cố định số lượng ứng viên $k$, ta cố định khối lượng xác suất tích lũy $p \in (0, 1]$ (thường chọn $p = 0.9$ hoặc $0.95$). Ta tìm tập con nhỏ nhất các token đứng đầu $V^{(p)}$ sao cho tổng xác suất của chúng vừa chạm hoặc vượt ngưỡng $p$:

   $$V^{(p)} = \left\{ w_{(1)}, \dots, w_{(m)} \;\middle|\; m = \min \left\{ j : \sum_{i=1}^{j} p_{(i)}(T) \ge p \right\} \right\}.$$

   Mọi token nằm ngoài $V^{(p)}$ đều bị loại bỏ, và các token trong hạt nhân được chuẩn hóa lại thành phân phối cuối cùng $\tilde{p}$:

   $$\tilde{p}(w) = \begin{cases} \dfrac{p(w)}{\sum_{v \in V^{(p)}} p(v)} & \text{nếu } w \in V^{(p)}, \\[8pt] 0 & \text{ngược lại.} \end{cases}$$

Nhờ định nghĩa theo tổng tích lũy, kích thước hạt nhân $|V^{(p)}|$ tự động co giãn linh hoạt: thu nhỏ về chỉ $1-2$ từ khi mô hình chắc chắn, và mở rộng ra hàng chục từ khi ngữ cảnh có nhiều nhánh rẽ hợp lý! Hãy thử nghiệm trực tiếp cơ chế này trên biểu đồ tương tác dưới đây.

Hãy tự tay cài đặt toàn bộ đường ống chọn mẫu này trong bài tập thực hành đầu tiên.

import torch

torch.manual_seed(42)
vocab = [
    "căn", "thế", "kho", "hành", "bóng",
    "khu", "người", "chiếc", "con", "trái",
]
logits = torch.tensor(
    [3.1, 2.7, 2.4, 2.0, 1.7, 1.3, 0.9, 0.5, 0.1, -0.6]
)
print("Từ vựng:", len(vocab), "tokens | Logit cao nhất:", logits[0].item())

Từ vựng: 10 tokens | Logit cao nhất: 3.0999999046325684

[Greedy T=0.1, k=1]  Giữ: 1/10 | Entropy: 0.00 bits | Top-1: căn (100.0%)
[Top-p  T=0.8, p=0.9] Giữ: 6/10 | Entropy: 2.26 bits | Top-1: căn (35.1%)
[Top-k  T=1.2, k=4]   Giữ: 4/10 | Entropy: 1.93 bits | Top-1: căn (37.1%)

Mẫu rút thử (T=0.8, k=8, p=0.90): căn, thế, căn, kho, hành, căn, thế, bóng, căn, kho

import torch

torch.manual_seed(42)
vocab = [
    "căn", "thế", "kho", "hành", "bóng",
    "khu", "người", "chiếc", "con", "trái",
]
logits = torch.tensor(
    [3.1, 2.7, 2.4, 2.0, 1.7, 1.3, 0.9, 0.5, 0.1, -0.6]
)
print("Từ vựng:", len(vocab), "tokens | Logit cao nhất:", logits[0].item())

Từ vựng: 10 tokens | Logit cao nhất: 3.0999999046325684

def filter_and_sample(logits, temp=1.0, top_k=10, top_p=1.0):
    # (1) Chia nhiệt độ T và sắp xếp logits giảm dần
    scaled = logits / max(temp, 1e-5)
    sorted_logits, sorted_idx = torch.sort(scaled, descending=True)

    # (2) Áp dụng mặt nạ Top-k
    k = min(top_k, sorted_logits.size(-1))
    sorted_logits[k:] = -float("inf")

    # (3) Tính xác suất tích lũy và áp dụng mặt nạ Top-p (Nucleus)
    probs = torch.softmax(sorted_logits, dim=-1)
    cum_probs = torch.cumsum(probs, dim=-1)
    remove_mask = (cum_probs - probs) >= top_p
    remove_mask[0] = False
    sorted_logits[remove_mask] = -float("inf")

    # (4) Chuẩn hóa lại xác suất trên thứ tự từ vựng gốc
    final_Sorted = torch.softmax(sorted_logits, dim=-1)
    final_probs = torch.zeros_like(logits)
    final_probs.scatter_(0, sorted_idx, final_Sorted)
    return final_probs


def calc_entropy(probs):
    nz = probs[probs > 0]
    return -(nz * torch.log2(nz)).sum().item()


configs = [
    ("Greedy T=0.1, k=1", 0.1, 1, 1.0),
    ("Top-p  T=0.8, p=0.9", 0.8, 10, 0.9),
    ("Top-k  T=1.2, k=4", 1.2, 4, 1.0),
]
for label, t, k, p in configs:
    dist = filter_and_sample(logits, temp=t, top_k=k, top_p=p)
    kept = int((dist > 0).sum().item())
    h = calc_entropy(dist)
    top_i = int(torch.argmax(dist).item())
    print(
        f"[{label:19s}] Giữ: {kept}/10 | Entropy: {h:.2f} bits"
        f" | Top-1: {vocab[top_i]} ({dist[top_i]*100:.1f}%)"
    )

torch.manual_seed(0)
dist_p = filter_and_sample(logits, temp=0.8, top_k=8, top_p=0.90)
draws = torch.multinomial(dist_p, num_samples=10, replacement=True)
words = ", ".join(vocab[i] for i in draws.tolist())
print(f"\nMẫu rút thử (T=0.8, k=8, p=0.90): {words}")

[Greedy T=0.1, k=1]  Giữ: 1/10 | Entropy: 0.00 bits | Top-1: căn (100.0%)
[Top-p  T=0.8, p=0.9] Giữ: 6/10 | Entropy: 2.26 bits | Top-1: căn (35.1%)
[Top-k  T=1.2, k=4]   Giữ: 4/10 | Entropy: 1.93 bits | Top-1: căn (37.1%)

Mẫu rút thử (T=0.8, k=8, p=0.90): căn, thế, căn, kho, hành, căn, thế, bóng, căn, kho

2. Bộ nhớ đệm KV (KV Cache) và giới hạn phần cứng Roofline

Sau khi đã biết cách rút ra token tiếp theo $x_{t+1}$, ta nối nó vào cuối chuỗi đầu vào để thành $x_1, \dots, x_t, x_{t+1}$ rồi lại đưa vào mô hình để sinh tiếp $x_{t+2}$. Nhưng nếu thực hiện một cách ngây thơ — mỗi bước lại đưa toàn bộ chuỗi từ đầu vào mạng Transformer — tốc độ suy luận sẽ sụp đổ rất nhanh khi câu trả lời dài ra. Tại sao lại như vậy?

2.1 Tại sao tính lại toàn bộ chuỗi tốn $O(L^2)$ và cách KV Cache giải cứu

Hãy nhìn vào một lớp Causal Self-Attention (Chương 2, mục 1) khi xử lý chuỗi có độ dài $t$ với ma trận biểu diễn đầu vào $X_{1:t} \in \mathbb{R}^{t \times d}$. Lớp này nhân $X_{1:t}$ với ba ma trận trọng số $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$ để tạo ra:

Có hai sự thật toán học vô cùng quan trọng ở đây:

1. Nhờ mặt nạ nhân quả (causal mask), biểu diễn của các token ở vị trí $1, 2, \dots, t-1$ hoàn toàn không phụ thuộc vào token mới $x_t$. Do đó, các hàng $k_1 \dots k_{t-1}$ và $v_1 \dots v_{t-1}$ ở bước $t$ giống hệt 100% với những gì chúng ta đã tính ở bước $t-1$!
2. Để dự đoán token tiếp theo $x_{t+1}$, lớp đầu ra cuối cùng của mô hình chỉ đọc duy nhất vector ẩn ở vị trí cuối cùng $h_t$. Mà để tính $h_t$, phép chú ý chỉ cần lấy tích vô hướng giữa truy vấn mới nhất $q_t \in \mathbb{R}^{1 \times d}$ với tất cả các khóa $k_1 \dots k_t$:

Nếu không lưu lại gì từ bước trước, ở mỗi bước $t$ ta phải tính lại toàn bộ $t$ hàng của $Q, K, V$; tổng chi phí sinh $L$ token sẽ tỷ lệ với $\sum_{t=1}^{L} t = \frac{L(L+1)}{2} = O(L^2)$.

Giải pháp rất tự nhiên: chúng ta cấp phát một vùng nhớ trên RAM của GPU gọi là Bộ nhớ đệm Khóa–Giá trị (KV Cache) để lưu lại các ma trận $K_{1:t-1}$ và $V_{1:t-1}$ ở mọi lớp Transformer. Khi sinh token thứ $t$, ta chỉ đưa duy nhất 1 token mới nhất $x_t$ đi qua mạng:

• Chỉ tính $q_t = x_t W_Q$, $k_t = x_t W_K$, $v_t = x_t W_V$ (chi phí $O(1)$ theo độ dài chuỗi).
• Nối (append) $k_t$ và $v_t$ vào cuối bộ đệm để thu được $K_{1:t}$ và $V_{1:t}$.
• Nhân vector $q_t$ ($1 \times d_k$) với $K_{1:t}^\top$ ($d_k \times t$) và $V_{1:t}$ ($t \times d_v$).

2.2 Hai pha suy luận (Prefill vs Decode) và Mô hình Roofline

KV Cache chia quá trình suy luận của một yêu cầu người dùng thành hai pha có bản chất phần cứng đối lập nhau hoàn toàn:

• Pha 1 — Prefill (Khởi tạo ngữ cảnh): Khi người dùng vừa gửi một câu hỏi (prompt) dài $L_{\text{prompt}}$ token (ví dụ $512$ token), toàn bộ $512$ token này đã có sẵn cùng lúc. GPU xử lý tất cả chúng song song trong một lượt chạy duy nhất bằng phép nhân Ma trận – Ma trận (kích thước $[512, d] \times [d, d]$), vừa sinh ra token trả lời đầu tiên, vừa điền đầy KV Cache ban đầu.
• Pha 2 — Decode (Giải mã tự hồi quy): Từ token thứ hai trở đi, vì token sau phụ thuộc vào token trước nên mô hình buộc phải chạy tuần tự từng bước một. Nhờ KV Cache, đầu vào của mỗi bước chỉ có độ dài bằng $1$ (hoặc bằng kích thước batch $B$). Mọi phép nhân trọng số giờ đây thu nhỏ thành phép nhân Vector – Ma trận (kích thước $[1, d] \times [d, d]$).

Tại sao phép nhân Vector – Ma trận ở pha Decode lại là “cơn ác mộng” đối với GPU hiện đại? Để hiểu điều này, các kỹ sư hệ thống sử dụng Mô hình Roofline (Roofline Model).

Một GPU có hai thông số vật lý giới hạn:

1. Đỉnh tính toán ($\text{Peak FLOP/s}$): Số phép tính dấu phẩy động tối đa mỗi giây của các nhân Tensor Cores (ví dụ NVIDIA A100 đạt $312\text{ TFLOPS} = 312 \times 10^{12}\text{ FLOP/s}$ ở chuẩn FP16).
2. Băng thông bộ nhớ ($\text{Memory Bandwidth}$): Tốc độ truyền dữ liệu tối đa từ bộ nhớ chính HBM (VRAM) vào bộ nhớ đệm siêu tốc trên chip SRAM (ví dụ A100 đạt $2.0\text{ TB/s} = 2 \times 10^{12}\text{ Bytes/s}$).

Tỷ số giữa số phép toán số học thực hiện được trên mỗi byte dữ liệu phải đọc từ HBM được gọi là Cường độ số học (Arithmetic Intensity — $I$):

Trên GPU A100, để các nhân tính toán không phải ngồi chơi chờ dữ liệu, một tác vụ phải có cường độ số học vượt qua điểm gãy (Ridge Point):

Hãy tính cường độ số học của một lớp tuyến $Y = X W$ với ma trận trọng số $W \in \mathbb{R}^{d \times d}$ lưu ở chuẩn FP16 ($2\text{ bytes}$ mỗi tham số, tổng dung lượng $2 d^2\text{ bytes}$) khi xử lý $N$ token cùng lúc ($X \in \mathbb{R}^{N \times d}$):

• Số phép tính nhân–cộng của phép nhân ma trận là $2 N d^2\text{ FLOPs}$.
• Khi $N \ll d$, lượng byte cần đọc từ HBM bị chi phối bởi kích thước ma trận trọng số $2 d^2\text{ bytes}$.
• Suy ra cường độ số học xấp xỉ bằng đúng số lượng token xử lý song song:

Đây chính là kết luận quan trọng nhất về phần cứng suy luận: Ở pha Decode với kích thước batch nhỏ, tốc độ sinh token (tokens/giây) hoàn toàn không phụ thuộc vào việc GPU có bao nhiêu TFLOPS, mà bằng Băng thông bộ nhớ chia cho Dung lượng mô hình! Ví dụ, mô hình Llama-3 8B ở chuẩn FP16 nặng $16\text{ GB}$; để sinh 1 token, GPU A100 ($2000\text{ GB/s}$) bắt buộc phải đọc hết $16\text{ GB}$ trọng số này từ HBM lên chip, nên tốc độ tối đa lý thuyết khi $B=1$ là:

Hãy quan sát cả cơ chế ma trận của KV Cache lẫn vị trí của hai pha Prefill / Decode trên biểu đồ Roofline dưới đây.

Hãy kiểm chứng sự tương đương toán học tuyệt đối và mức giảm số phép tính của KV Cache trong Bài tập 5.2.

import math
import torch
import torch.nn as nn

torch.manual_seed(7)
B, L_prompt, gen_steps, d_model = 1, 8, 16, 64
W_q = nn.Linear(d_model, d_model, bias=False)
W_k = nn.Linear(d_model, d_model, bias=False)
W_v = nn.Linear(d_model, d_model, bias=False)
W_o = nn.Linear(d_model, d_model, bias=False)

seq_tokens = torch.randn(B, L_prompt + gen_steps, d_model)
print("Kích thước chuỗi mô phỏng:", tuple(seq_tokens.shape))

Kích thước chuỗi mô phỏng: (1, 24, 64)

Sai số tuyệt đối lớn nhất (No-Cache vs KV-Cache): 2.98e-08
Kích thước bộ đệm K, V sau cùng: (1, 24, 64)
Số lượt chiếu QKV tuyến tính (Không Cache): 268 lượt
Số lượt chiếu QKV tuyến tính (Có KV Cache): 24 lượt (Giảm 11.2x)

import math
import torch
import torch.nn as nn

torch.manual_seed(7)
B, L_prompt, gen_steps, d_model = 1, 8, 16, 64
W_q = nn.Linear(d_model, d_model, bias=False)
W_k = nn.Linear(d_model, d_model, bias=False)
W_v = nn.Linear(d_model, d_model, bias=False)
W_o = nn.Linear(d_model, d_model, bias=False)

seq_tokens = torch.randn(B, L_prompt + gen_steps, d_model)
print("Kích thước chuỗi mô phỏng:", tuple(seq_tokens.shape))

Kích thước chuỗi mô phỏng: (1, 24, 64)

def causal_attn_step(x, kv_cache=None):
    # (1) Chiếu tuyến tính Q, K, V cho các token đầu vào x
    q, k, v = W_q(x), W_k(x), W_v(x)
    qkv_calls = x.size(1)

    # (2) Nối k, v mới vào bộ nhớ đệm kv_cache nếu có
    if kv_cache is not None:
        k_prev, v_prev = kv_cache
        k = torch.cat([k_prev, k], dim=1)
        v = torch.cat([v_prev, v], dim=1)
    new_cache = (k, v)

    # (3) Tính điểm chú ý nhân quả và đầu ra cho vị trí cuối cùng
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(d_model)
    L_q, L_k = q.size(1), k.size(1)
    if L_q > 1:
        mask = torch.triu(
            torch.ones(L_q, L_k, dtype=torch.bool), diagonal=1
        )
        scores = scores.masked_fill(mask, -float("inf"))
    weights = torch.softmax(scores, dim=-1)
    out = W_o(weights @ v)
    return out[:, -1:, :], new_cache, qkv_calls


with torch.no_grad():
    # (A) Không dùng KV Cache: mỗi bước truyền lại toàn bộ tiền tố
    outs_nocache, calls_nocache = [], 0
    for step in range(gen_steps + 1):
        cur_len = L_prompt + step
        out_t, _, c = causal_attn_step(seq_tokens[:, :cur_len, :])
        outs_nocache.append(out_t)
        calls_nocache += c

    # (B) Có KV Cache: pha Prefill (8 token) + 16 bước Decode (1 token)
    outs_cached, calls_cached = [], 0
    out_0, cache, c = causal_attn_step(seq_tokens[:, :L_prompt, :])
    outs_cached.append(out_0)
    calls_cached += c
    for step in range(1, gen_steps + 1):
        tok_t = seq_tokens[:, L_prompt + step - 1 : L_prompt + step, :]
        out_t, cache, c = causal_attn_step(tok_t, kv_cache=cache)
        outs_cached.append(out_t)
        calls_cached += c

diff = (torch.cat(outs_nocache, 1) - torch.cat(outs_cached, 1)).abs()
print(f"Sai số tuyệt đối lớn nhất (No-Cache vs KV-Cache): {diff.max():.2e}")
print("Kích thước bộ đệm K, V sau cùng:", tuple(cache[0].shape))
print(f"Số lượt chiếu QKV tuyến tính (Không Cache): {calls_nocache} lượt")
ratio = calls_nocache / calls_cached
print(
    f"Số lượt chiếu QKV tuyến tính (Có KV Cache): {calls_cached} lượt"
    f" (Giảm {ratio:.1f}x)"
)

Sai số tuyệt đối lớn nhất (No-Cache vs KV-Cache): 2.98e-08
Kích thước bộ đệm K, V sau cùng: (1, 24, 64)
Số lượt chiếu QKV tuyến tính (Không Cache): 268 lượt
Số lượt chiếu QKV tuyến tính (Có KV Cache): 24 lượt (Giảm 11.2x)

3. Thu gọn KV Cache: MQA, GQA và PagedAttention

KV Cache giúp chúng ta thoát khỏi độ phức tạp tính toán bậc hai $O(L^2)$ ở mỗi bước Decode, nhưng nó lại đánh đổi bằng một cái giá rất đắt: dung lượng bộ nhớ VRAM tăng tuyến tính theo cả độ dài ngữ cảnh $L$ lẫn kích thước batch $B$.

Hãy làm một phép tính nhẩm cho mô hình Llama-2 70B nếu dùng cơ chế Multi-Head Attention (MHA) tiêu chuẩn ($n_{\text{layers}} = 80$ lớp, $n_{\text{heads}} = 64$ đầu chú ý, chiều mỗi đầu $d_h = 128$, lưu ở chuẩn FP16 tốn $2\text{ bytes}$ cho mỗi số thực):

• Với mỗi token, ở mỗi lớp và mỗi đầu chú ý, ta phải lưu 1 vector khóa $k$ ($128$ số) và 1 vector giá trị $v$ ($128$ số).
• Dung lượng KV Cache cho 1 token duy nhất trên toàn mạng là: $$2\text{ (K và V)} \times 2\text{ bytes} \times 80\text{ lớp} \times 64\text{ đầu} \times 128\text{ chiều} = 2{,}621{,}440\text{ bytes} \approx 2.5\text{ MB / token}.$$
• Nếu hệ thống phục vụ đồng thời một batch gồm $B = 16$ người dùng với độ dài ngữ cảnh $L = 8{,}192$ token, tổng số token cần lưu trong KV Cache là $16 \times 8{,}192 = 131{,}072$ token. Dung lượng VRAM riêng cho KV Cache sẽ là: $$131{,}072 \times 2.5\text{ MB} = \mathbf{320\text{ GB VRAM!}}$$

Con số này còn lớn gấp đôi dung lượng của bản thân trọng số mô hình 70B ($140\text{ GB}$)! Hơn thế nữa, ở mỗi bước Decode, GPU phải đọc toàn bộ $320\text{ GB}$ KV Cache này từ HBM lên chip để tính tích vô hướng $q_t K^\top$, khiến nút cổ chai băng thông bộ nhớ càng trở nên trầm trọng.

3.1 Từ Multi-Query Attention (MQA) đến Grouped-Query Attention (GQA)

Làm thế nào để cắt giảm kích thước KV Cache đi nhiều lần mà không làm giảm số lượng đầu truy vấn (Query heads) của mô hình?

Năm 2019, Noam Shazeer (một trong các đồng tác giả của Transformer gốc) nhận ra một sự bất đối xứng thú vị (Shazeer, 2019): trong Multi-Head Attention, chúng ta rất cần nhiều đầu truy vấn $Q^{(1)}, \dots, Q^{(H)}$ để mỗi đầu đặt một “câu hỏi” khác nhau ở mỗi bước. Nhưng các bảng tra cứu $K$ và $V$ thì hoàn toàn có thể dùng chung! Ông đề xuất Multi-Query Attention (MQA): giữ nguyên toàn bộ $H$ đầu Query, nhưng chỉ sử dụng duy nhất 1 đầu Key và 1 đầu Value chia sẻ cho tất cả $H$ đầu Query đó. Với $H = 64$, MQA lập tức làm giảm dung lượng và băng thông đọc của KV Cache đi đúng 64 lần!

Tuy nhiên, việc ép cả 64 đầu Query phải nhìn vào cùng 1 không gian Key/Value duy nhất đôi khi làm giảm độ chính xác trên các bài toán suy luận phức tạp. Năm 2023, Ainslie và cộng sự tại Google đã đưa ra điểm cân bằng vàng mang tên Grouped-Query Attention (GQA) (Ainslie & others, 2023):

• Ta chia $n_q$ đầu Query thành $G$ nhóm bằng nhau (với $1 < G < n_q$, tức có $n_{\text{kv}} = G$ đầu Key/Value).
• Mỗi nhóm gồm $n_q / G$ đầu Query sẽ chia sẻ chung 1 đầu Key và 1 đầu Value.
• Công thức tổng quát tính dung lượng bộ nhớ KV Cache (tính bằng Bytes) cho batch size $B$ và độ dài chuỗi $L$ giờ đây là:

Ngày nay, GQA với $n_{\text{kv}} = 8$ đầu Key/Value đã trở thành tiêu chuẩn mặc định của hầu hết mọi mô hình ngôn ngữ lớn hiện đại (bao gồm Llama-3 8B/70B, Qwen-2.5, Mistral và Gemma-2), giúp giảm dung lượng KV Cache đi $4\times$ đến $8\times$ hoàn toàn miễn phí về mặt chất lượng!

3.2 Quản lý bộ nhớ phân trang PagedAttention (vLLM)

Ngay cả khi đã dùng GQA để thu gọn kích thước mỗi token trong KV Cache, một vấn đề hệ thống cực lớn khác vẫn tồn tại khi vận hành máy chủ suy luận cho nhiều người dùng cùng lúc: sự phân mảnh bộ nhớ (memory fragmentation).

Trong PyTorch thông thường, một tensor kv_cache của một yêu cầu phải được cấp phát dưới dạng một khối bộ nhớ liên tục (contiguous memory) trên VRAM. Nhưng khi người dùng vừa gửi câu hỏi vào, hệ thống không thể biết trước câu trả lời sẽ dài $50$ token hay $2{,}048$ token! Do đó, các hệ thống thế hệ đầu buộc phải đặt chỗ trước (pre-allocate) một vùng nhớ liên tục bằng đúng độ dài tối đa $L_{\text{max}}$ (ví dụ $4{,}096$ token) cho mỗi yêu cầu.

Cách cấp phát liên tục này gây ra sự lãng phí khủng khiếp (Kwon & others, 2023):

1. Phân mảnh nội (Internal fragmentation): Vùng nhớ đã đặt chỗ cho $4{,}096$ token nhưng thực tế câu trả lời chỉ dài $300$ token; hơn $90\%$ phần bộ nhớ còn lại bị bỏ không mà không yêu cầu nào khác được chạm vào.
2. Phân mảnh ngoại (External fragmentation): Khi các yêu cầu có độ dài khác nhau kết thúc và giải phóng bộ nhớ, VRAM bị băm nhỏ thành những khoảng trống rời rạc không đủ lớn để chứa một khối liên tục mới.

Năm 2023, nhóm nghiên cứu tại UC Berkeley đã giải quyết triệt để bài toán này bằng PagedAttention (vLLM) (Kwon & others, 2023), mượn trực tiếp nguyên lý bộ nhớ ảo phân trang (virtual memory paging) của hệ điều hành Linux/Windows:

• Thay vì cấp phát một mảng lớn liên tục cho mỗi yêu cầu, toàn bộ VRAM dành cho KV Cache được chia thành các khối (KV blocks / pages) có kích thước nhỏ cố định (ví dụ mỗi khối chứa đúng $B_{\text{block}} = 16$ token).
• Một yêu cầu sinh đến đâu mới được cấp phát thêm khối mới đến đó (on-demand allocation). Các khối của cùng một cuộc hội thoại không cần nằm cạnh nhau trong bộ nhớ vật lý HBM, mà được ánh xạ thông qua một Bảng khối (Block Table).
• Nhân tính toán CUDA của PagedAttention đọc trực tiếp bảng khối này để gom các vector $K, V$ từ các trang vật lý rời rạc ngay trong lúc tính điểm chú ý.

Đặc biệt, PagedAttention còn mở ra khả năng chia sẻ tiền tố (Prefix Caching / Copy-on-Write): khi hàng nghìn yêu cầu cùng dùng chung một đoạn System Prompt dài (hoặc tài liệu RAG chung), tất cả chúng chỉ cần trỏ chung vào cùng các trang KV vật lý của đoạn mở đầu đó mà không tốn thêm một byte VRAM nào!

4. FlashAttention và lượng tử hóa trọng số (Quantization)

Ở Mục 2 và Mục 3, chúng ta đã thấy bức tường băng thông bộ nhớ (Memory Wall) kìm hãm cả hai thành phần lớn nhất của mô hình:

1. Ở pha Prefill (hoặc khi ngữ cảnh $L$ rất dài), việc đọc/ghi ma trận chú ý trung gian $S, P \in \mathbb{R}^{L \times L}$ ra vào bộ nhớ HBM làm tắc nghẽn cơ chế Self-Attention.
2. Ở pha Decode, việc phải đọc toàn bộ ma trận trọng số $W$ từ HBM cho mỗi token làm tắc nghẽn các lớp tuyến tính (Linear / FFN).

Chúng ta sẽ lần lượt phá vỡ hai nút cổ chai này bằng FlashAttention và Lượng tử hóa Trọng số (Weight Quantization).

4.1 FlashAttention: Tránh ghi ma trận $L \times L$ ra HBM nhờ Online Softmax

Trên một GPU như A100, bộ nhớ được tổ chức theo cấu trúc phân cấp với sự chênh lệch tốc độ rất lớn:

• Bộ nhớ chính HBM (VRAM): Dung lượng rất lớn ($40\text{ GB} - 80\text{ GB}$), nhưng băng thông chậm ($1.5 - 2.0\text{ TB/s}$).
• Bộ nhớ đệm trên chip SRAM (Shared Memory / L1 Cache): Băng thông cực nhanh ($\approx 19\text{ TB/s}$, nhanh gấp $10\times$ HBM), nhưng dung lượng cực nhỏ (chỉ khoảng $192\text{ KB}$ trên mỗi Streaming Multiprocessor, tổng cộng $\approx 20\text{ MB}$ trên toàn chip).

Cách cài đặt Self-Attention tiêu chuẩn trong PyTorch ($\text{softmax}(QK^\top / \sqrt{d_k}) V$) thực hiện 3 bước tách rời, buộc GPU phải ghi các ma trận trung gian khổng lồ kích thước $L \times L$ ra HBM rồi lại đọc ngược vào SRAM:

1. Đọc $Q, K \in \mathbb{R}^{L \times d}$ từ HBM lên SRAM, tính $S = QK^\top / \sqrt{d_k} \in \mathbb{R}^{L \times L}$, rồi ghi toàn bộ $S$ ra HBM.
2. Đọc $S$ từ HBM lên SRAM, tính $P = \text{softmax}(S) \in \mathbb{R}^{L \times L}$, rồi ghi toàn bộ $P$ ra HBM.
3. Đọc $P$ và $V$ từ HBM lên SRAM, tính $O = PV \in \mathbb{R}^{L \times d}$, rồi ghi $O$ ra HBM.

Khi độ dài chuỗi $L = 32{,}768$, ma trận $L \times L$ có tới hơn 1 tỷ phần tử cho mỗi đầu chú ý! Thời gian di chuyển ma trận $L \times L$ này qua lại giữa HBM và SRAM lâu gấp nhiều lần thời gian thực hiện phép nhân ma trận thực sự.

Năm 2022, Tri Dao và cộng sự giới thiệu FlashAttention (Dao & others, 2022) với ý tưởng cốt lõi: Không bao giờ tạo ra và ghi ma trận $L \times L$ xuống HBM! Thay vào đó, ta chia (tile) các ma trận $Q, K, V$ thành từng khối nhỏ nạp vừa vào SRAM, tính điểm chú ý của từng khối rồi cộng dồn trực tiếp vào kết quả đầu ra $O$ ngay bên trong SRAM trước khi ghi $O$ về HBM.

Nhưng có một rào cản toán học: hàm $\text{softmax}(s_i) = \frac{e^{s_i - m}}{\sum_{j=1}^{L} e^{s_j - m}}$ (với $m = \max_j s_j$ để tránh tràn số) có mẫu số phụ thuộc vào toàn bộ hàng $1 \dots L$. Làm sao ta có thể nhân $P$ với $V$ khi mới chỉ đọc được khối thứ nhất của $K$ và $V$?

Bí quyết nằm ở kỹ thuật Online Softmax (chuẩn hóa lại trực tuyến). Giả sử khi đọc xong khối thứ nhất $S^{(1)}$, ta lưu giá trị lớn nhất cục bộ $m^{(1)} = \max(S^{(1)})$, tổng mẫu số cục bộ $\ell^{(1)} = \sum e^{S^{(1)} - m^{(1)}}$, và vector đầu ra tạm thời $O^{(1)} = \frac{e^{S^{(1)} - m^{(1)}}}{\ell^{(1)}} V^{(1)}$. Khi nạp tiếp khối thứ hai $S^{(2)}, V^{(2)}$ vào SRAM, ta cập nhật thống kê toàn cục mới chỉ bằng một phép nhân hiệu chỉnh hệ số mũ:

Nhờ quy tắc cập nhật đệ quy tuyệt đẹp này, GPU có thể quét qua từng mảnh của $K, V$ trong SRAM và cuối cùng thu được chính xác tuyệt đối ma trận $O$ mà không hề tốn một byte HBM nào cho ma trận $L \times L$!

4.2 Lượng tử hóa trọng số (Weight Quantization): Nén FP16 xuống INT8 và INT4

Nếu FlashAttention giải quyết nút cổ chai $L \times L$ của cơ chế Attention, thì ở pha Decode (với độ dài chuỗi vừa phải), hơn $80\%$ thời gian lại nằm ở việc đọc các ma trận trọng số $W \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$ từ HBM vào SRAM.

Vì tốc độ pha Decode tỷ lệ nghịch với dung lượng byte của trọng số:

• Nếu ta nén mỗi trọng số từ chuẩn 16-bit (FP16/BF16, $2\text{ bytes}$) xuống số nguyên 8-bit (INT8, $1\text{ byte}$), dung lượng mô hình giảm $2\times$ và tốc độ sinh token lý thuyết tăng gấp $2\times$.
• Nếu nén xuống số nguyên 4-bit (INT4, $0.5\text{ byte}$), dung lượng mô hình 70B giảm từ $140\text{ GB}$ xuống chỉ còn $\approx 38\text{ GB}$ (chạy vừa trên một card GPU tiêu dùng hoặc Mac Studio!), và tốc độ đọc bộ nhớ tăng gần $4\times$.

Cơ chế lượng tử hóa hậu huấn luyện (Post-Training Quantization — PTQ) phổ biến nhất là lượng tử hóa đối xứng Absmax (Symmetric Absmax Quantization). Để ánh xạ một vector thực $\mathbf{w} \in \mathbb{R}^g$ sang lưới số nguyên $b$-bit nằm trong đoạn $[-q_{\max}, q_{\max}]$ với $q_{\max} = 2^{b-1} - 1$ (ví dụ $q_{\max} = 127$ cho INT8 và $q_{\max} = 7$ cho INT4), ta tìm giá trị tuyệt đối lớn nhất và định nghĩa hệ số tỷ lệ (scale factor) $s$:

Khi thực hiện phép nhân ma trận trong lúc suy luận, các trọng số nguyên $w_i^{(q)}$ được nạp từ HBM lên SRAM với băng thông cực nhỏ, sau đó được giải lượng tử hóa (dequantize) ngay trên chip về lại số thực xấp xỉ $\hat{w}_i = s \cdot w_i^{(q)}$ để nhân với vector kích hoạt đầu vào $x$.

Tuy nhiên, nếu dùng chung một hệ số tỷ lệ $s$ cho toàn bộ ma trận (Per-Tensor Quantization), một rắc rối lớn sẽ xuất hiện trên các LLM có quy mô từ 6.7 tỷ tham số trở lên (Dettmers & others, 2022): các đặc trưng ngoại lai (emergent outlier features). Một số kênh ẩn có độ lớn vượt trội gấp $10\times - 20\times$ các kênh bình thường. Khi một phần tử ngoại lai kéo $\max |w_i|$ lên rất cao, hệ số tỷ lệ $s$ trở nên quá lớn, khiến $99.9\%$ các trọng số nhỏ còn lại trong ma trận bị làm tròn hết về con số $0$!

Để khắc phục hiện tượng này, bốn phương pháp kinh điển đã lần lượt ra đời:

• Chia nhóm (Group-wise Quantization): Thay vì lấy một hệ số $s$ cho cả ma trận hay cả hàng, ta chia mỗi hàng trọng số thành các nhóm nhỏ liên tiếp có kích thước $g \in \{32, 64, 128\}$ phần tử. Mỗi nhóm có một hệ số $s$ (FP16) riêng biệt. Một phần tử ngoại lai giờ đây chỉ ảnh hưởng tới 64 phần tử trong nhóm của nó thay vì làm hỏng cả ma trận, trong khi chi phí lưu thêm các hệ số $s$ chỉ tốn thêm $2\text{ bytes} / 64 = 0.25\text{ bits}$ mỗi tham số!
• LLM.int8() [@dettmers2022llmint8]: Phát hiện ra rằng các giá trị ngoại lai luôn tập trung ở một vài chiều kênh cố định (outlier dimensions). LLM.int8() tách riêng $<0.1\%$ các cột ngoại lai này ra để nhân ở chuẩn FP16, còn $99.9\%$ ma trận còn lại được nhân bằng nhân phần cứng INT8.
• GPTQ [@frantar2023gptq]: Lượng tử hóa từng cột trọng số xuống 4-bit và sử dụng thông tin đạo hàm bậc hai (nghịch đảo ma trận Hessian $H^{-1} = (2 X X^\top)^{-1}$ tính trên một tập dữ liệu hiệu chuẩn nhỏ) để cập nhật bù trừ sai số làm tròn của cột vừa lượng tử hóa sang các cột chưa lượng tử hóa phía sau.
• AWQ (Activation-aware Weight Quantization) [@lin2024awq]: Nhận thấy rằng không phải mọi trọng số đều quan trọng như nhau — những kênh trọng số đi kèm với đầu vào kích hoạt $X$ có biên độ lớn mới là nơi quyết định sai số của tích $X W^\top$. AWQ nhân phóng đại (per-channel scaling) $1\%$ các kênh trọng số nhạy cảm nhất trước khi lượng tử hóa 4-bit để bảo vệ chúng khỏi sai số làm tròn mà không cần dùng độ chính xác hỗn hợp.

Hãy tự tay cài đặt lượng tử hóa theo nhóm (Group-wise INT8 & INT4) và kiểm chứng tác động của các kênh ngoại lai trong Bài tập 5.4 dưới đây.

import torch

torch.manual_seed(123)
d_out, d_in = 256, 512
W = torch.randn(d_out, d_in, dtype=torch.float32) * 0.08
# Chèn một vài kênh ngoại lai (outlier channels) như trong LLM thực tế
W[:, 17] *= 8.5
W[:, 203] *= 6.0

X = torch.randn(16, d_in, dtype=torch.float32)
Y_ref = X @ W.T
print(f"Ma trận W: {tuple(W.shape)} | Dung lượng FP16 gốc: {W.numel()*2/1024:.1f} KB")

Ma trận W: (256, 512) | Dung lượng FP16 gốc: 256.0 KB

[INT8 Per-Tensor      ] Bộ nhớ: 128.0 KB (Nén 2.00x) | Sai số Y: 1.84%
[INT8 Group-wise g=64 ] Bộ nhớ: 132.0 KB (Nén 1.94x) | Sai số Y: 0.49%
[INT4 Per-Tensor      ] Bộ nhớ:  64.0 KB (Nén 4.00x) | Sai số Y: 31.25%
[INT4 Group-wise g=64 ] Bộ nhớ:  68.0 KB (Nén 3.76x) | Sai số Y: 8.42%

import torch

torch.manual_seed(123)
d_out, d_in = 256, 512
W = torch.randn(d_out, d_in, dtype=torch.float32) * 0.08
# Chèn một vài kênh ngoại lai (outlier channels) như trong LLM thực tế
W[:, 17] *= 8.5
W[:, 203] *= 6.0

X = torch.randn(16, d_in, dtype=torch.float32)
Y_ref = X @ W.T
print(f"Ma trận W: {tuple(W.shape)} | Dung lượng FP16 gốc: {W.numel()*2/1024:.1f} KB")

Ma trận W: (256, 512) | Dung lượng FP16 gốc: 256.0 KB

def quantize_dequantize(W, bits=8, group_size=None):
    q_max = (1 << (bits - 1)) - 1
    orig_shape = W.shape
    # (1) Chia ma trận thành từng nhóm kích thước group_size
    if group_size is None:
        W_g = W.reshape(1, -1)
    else:
        W_g = W.reshape(-1, group_size)

    # (2) Tính hệ số tỷ lệ Absmax cho từng nhóm và làm tròn
    abs_max = W_g.abs().amax(dim=1, keepdim=True).clamp(min=1e-8)
    scales = abs_max / q_max
    W_q = torch.clamp(torch.round(W_g / scales), -q_max, q_max)

    # (3) Giải lượng tử hóa (Dequantize) và tính dung lượng bộ nhớ
    W_hat = (W_q * scales).reshape(orig_shape)
    weight_bytes = W.numel() * (bits / 8.0)
    scale_bytes = scales.numel() * 2.0  # Lưu scale ở dạng FP16
    total_kb = (weight_bytes + scale_bytes) / 1024.0
    return W_hat, total_kb


fp16_kb = (W.numel() * 2.0) / 1024.0
configs = [
    ("INT8 Per-Tensor", 8, None),
    ("INT8 Group-wise g=64", 8, 64),
    ("INT4 Per-Tensor", 4, None),
    ("INT4 Group-wise g=64", 4, 64),
]

for name, bits, g_size in configs:
    W_hat, mem_kb = quantize_dequantize(W, bits=bits, group_size=g_size)
    Y_hat = X @ W_hat.T
    rel_err = (torch.norm(Y_ref - Y_hat) / torch.norm(Y_ref)).item() * 100
    ratio = fp16_kb / mem_kb
    print(
        f"[{name:21s}] Bộ nhớ: {mem_kb:5.1f} KB (Nén {ratio:.2f}x)"
        f" | Sai số Y: {rel_err:.2f}%"
    )

[INT8 Per-Tensor      ] Bộ nhớ: 128.0 KB (Nén 2.00x) | Sai số Y: 1.84%
[INT8 Group-wise g=64 ] Bộ nhớ: 132.0 KB (Nén 1.94x) | Sai số Y: 0.49%
[INT4 Per-Tensor      ] Bộ nhớ:  64.0 KB (Nén 4.00x) | Sai số Y: 31.25%
[INT4 Group-wise g=64 ] Bộ nhớ:  68.0 KB (Nén 3.76x) | Sai số Y: 8.42%

5. Giải mã suy đoán (Speculative Decoding)

Ngay cả khi chúng ta đã trang bị KV Cache, GQA, FlashAttention và lượng tử hóa trọng số, bản chất tự hồi quy của pha Decode vẫn giữ nguyên một định luật khắc nghiệt: muốn sinh ra $K$ token từ một mô hình đích khổng lồ (Target Model, ví dụ 70B), ta bắt buộc phải đọc toàn bộ trọng số của mô hình 70B đó từ HBM lên chip đúng $K$ lần tuần tự.

Trong khi đó, như ta đã thấy trên biểu đồ Roofline ở Mục 2.2, khi chạy với kích thước batch nhỏ ($B=1$), cường độ số học của GPU cực thấp: thời gian để mô hình 70B xử lý 1 token và thời gian để nó xử lý song song 5 token cùng lúc là gần như bằng nhau (vì cả hai đều chỉ đọc trọng số 70B từ HBM đúng 1 lần!).

Câu hỏi táo bạo được đặt ra là: Liệu ta có thể kiểm tra song song 5 token cùng lúc trong 1 lần chạy mô hình lớn, khi mà ta còn chưa biết 5 token đó là gì?

5.1 Nguyên lý Đề xuất và Kiểm chứng (Draft-and-Verify)

Năm 2023, Yaniv Leviathan và các cộng sự tại Google Research đã đưa ra câu trả lời thông qua thuật toán Giải mã Suy đoán (Speculative Decoding) (Leviathan & Matias, 2023). Ý tưởng xuất phát từ một quan sát rất thực tế: trong bất kỳ đoạn văn bản nào, phần lớn các token tiếp theo đều là những từ nối, dấu câu hoặc cụm từ cố định rất dễ đoán (như "Hà" $\to$ "Nội", "Việt" $\to$ "Nam", "import" $\to$ "torch"); chỉ có một số ít vị trí mới thực sự cần đến trí thông minh của mô hình 70B.

Thay vì dùng mô hình đích lớn $M_p$ (phân phối $p(x \mid x_{<t})$) để mò mẫm từng bước, ta đặt cạnh nó một mô hình nháp nhỏ (Draft Model — $M_q$) có phân phối $q(x \mid x_{<t})$ (ví dụ mô hình 1B cùng họ, nhẹ hơn và nhanh hơn gấp $15\times - 20\times$). Mỗi chu kỳ suy đoán gồm hai bước:

1. Bước Đề xuất (Drafting): Chạy tuần tự mô hình nháp nhỏ $M_q$ trong $\gamma$ bước (thường chọn $\gamma = 4$ hoặc $5$) để sinh nhanh một chuỗi $\gamma$ token ứng viên $\tilde{x}_1, \tilde{x}_2, \dots, \tilde{x}_\gamma$ cùng các xác suất nháp tương ứng $q_1(\tilde{x}_1), \dots, q_\gamma(\tilde{x}_\gamma)$.
2. Bước Kiểm chứng song song (Parallel Verification): Nối cả $\gamma$ token nháp này vào sau ngữ cảnh hiện tại và đưa qua mô hình đích lớn $M_p$ chỉ trong 1 lượt chạy tiến (forward pass) duy nhất. Nhờ tính chất song song của Transformer (giống hệt pha Prefill!), chỉ với 1 lần đọc trọng số từ HBM, mô hình lớn $M_p$ trả về đồng thời tất cả $\gamma + 1$ phân phối xác suất đích: $p_1(x), p_2(x), \dots, p_\gamma(x)$ và cả phân phối ở vị trí tiếp theo $p_{\gamma+1}(x)$!

5.2 Thuật toán Lấy mẫu Bác bỏ Suy đoán (Speculative Rejection Sampling)

Nếu chúng ta đang dùng giải mã tham lam ($T = 0$), việc kiểm chứng rất đơn giản: ta duyệt từ $i = 1$ đến $\gamma$, giữ lại các token nháp chừng nào $\tilde{x}_i == \arg\max p_i(x)$, và dừng lại ngay tại vị trí đầu tiên mà mô hình nháp đoán khác mô hình đích.

Nhưng khi chọn mẫu ngẫu nhiên ($T > 0$), làm thế nào để chấp nhận hoặc bác bỏ các token nháp $\tilde{x}_i \sim q_i(x)$ sao cho token đầu ra cuối cùng tuân theo đúng 100% phân phối đích $p_i(x)$ mà không hề bị thiên lệch theo mô hình nháp?

Leviathan et al. (Leviathan & Matias, 2023) giải quyết bài toán này bằng Lấy mẫu Bác bỏ Suy đoán (Speculative Rejection Sampling). Duyệt tuần tự từng vị trí $i = 1, \dots, \gamma$:

• Rút một số ngẫu nhiên phân phối đều $u \sim \text{Uniform}(0, 1)$ và so sánh với tỷ số xác suất giữa mô hình đích và mô hình nháp tại token ứng viên $\tilde{x}_i$: $$\text{Xác suất chấp nhận } \tilde{x}_i = \min\!\left(1,\, \frac{p_i(\tilde{x}_i)}{q_i(\tilde{x}_i)}\right).$$
• Trường hợp 1 (Chấp nhận — $u < \frac{p_i(\tilde{x}_i)}{q_i(\tilde{x}_i)}$): Token $\tilde{x}_i$ được giữ lại chính thức! Ta chuyển sang kiểm tra vị trí $i + 1$. (Lưu ý: nếu mô hình đích đánh giá $\tilde{x}_i$ cao hơn hoặc bằng mô hình nháp, tức $p_i(\tilde{x}_i) \ge q_i(\tilde{x}_i)$, thì token đó luôn luôn được chấp nhận $100\%$).
• Trường hợp 2 (Bác bỏ — $u \ge \frac{p_i(\tilde{x}_i)}{q_i(\tilde{x}_i)}$): Ta bác bỏ $\tilde{x}_i$ cùng toàn bộ các token nháp phía sau nó ($i+1 \dots \gamma$). Tuy nhiên, ta không phải chạy lại mô hình lớn! Ta lập tức rút một token thay thế $x_i'$ tại vị trí $i$ từ phân phối phần dư đã chuẩn hóa (normalized residual distribution) $p'_i(x)$ chỉ chứa những token mà mô hình nháp đã đánh giá thấp hơn mô hình đích: $$p'_i(x) = \frac{\max\!\big(0,\, p_i(x) - q_i(x)\big)}{\sum_{w \in V} \max\!\big(0,\, p_i(w) - q_i(w)\big)},$$ rồi kết thúc chu kỳ hiện tại.
• Trường hợp Thưởng đặc biệt: Nếu toàn bộ $\gamma$ token nháp đều được chấp nhận, ta lấy mẫu thêm 1 token thứ $\gamma + 1$ hoàn toàn miễn phí từ phân phối $p_{\gamma+1}(x)$ mà mô hình lớn đã tính sẵn ở đầu ra cuối cùng! Như vậy, mỗi chu kỳ luôn sinh ra ít nhất $1$ token và tối đa $\gamma + 1$ token mới.

Nếu gọi $\alpha \in (0, 1)$ là tỷ lệ chấp nhận trung bình mỗi bước và $c = T_{\text{draft}} / T_{\text{target}}$ là tỷ số thời gian chạy 1 bước của mô hình nháp so với mô hình đích, thì kỳ vọng số token sinh ra trong một chu kỳ là $\mathbb{E}[\text{tokens}] = \frac{1 - \alpha^{\gamma + 1}}{1 - \alpha}$, và hệ số tăng tốc thực tế là:

Hãy khám phá trực quan từng bước chấp nhận/bác bỏ và đường cong tăng tốc trên biểu đồ tương tác dưới đây.

Hãy kết thúc chương học bằng việc tự tay cài đặt bước kiểm chứng suy đoán và xác nhận định lý bảo toàn phân phối bằng mô phỏng Monte Carlo trong Bài tập 5.5.

import torch

torch.manual_seed(2026)
tokens = ["Hà", "Nội", "Sài", "Gòn", "Đà", "Nẵng"]
# Phân phối của mô hình đích lớn p(x) và mô hình nháp nhỏ q(x)
p_target = torch.tensor([0.42, 0.25, 0.15, 0.10, 0.05, 0.03])
q_draft = torch.tensor([0.22, 0.38, 0.18, 0.08, 0.09, 0.05])

print("Tổng p(x):", p_target.sum().item(), "| Tổng q(x):", q_draft.sum().item())

Tổng p(x): 1.0 | Tổng q(x): 1.0

Tỷ lệ chấp nhận lý thuyết α = Σ min(p, q): 78.00%
Tỷ lệ chấp nhận thực nghiệm (200,000 mẫu):  77.96%
Sai lệch cực đại |p_spec(x) - p(x)|:       0.0014

Token  |  p_target(x) |  q_draft(x)  |  p_spec(x)
-------------------------------------------------
Hà     |     0.4200   |    0.2200    |    0.4191
Nội    |     0.2500   |    0.3800    |    0.2514
Sài    |     0.1500   |    0.1800    |    0.1498
Gòn    |     0.1000   |    0.0800    |    0.0999
Đà     |     0.0500   |    0.0900    |    0.0497
Nẵng   |     0.0300   |    0.0500    |    0.0301

import torch

torch.manual_seed(2026)
tokens = ["Hà", "Nội", "Sài", "Gòn", "Đà", "Nẵng"]
# Phân phối của mô hình đích lớn p(x) và mô hình nháp nhỏ q(x)
p_target = torch.tensor([0.42, 0.25, 0.15, 0.10, 0.05, 0.03])
q_draft = torch.tensor([0.22, 0.38, 0.18, 0.08, 0.09, 0.05])

print("Tổng p(x):", p_target.sum().item(), "| Tổng q(x):", q_draft.sum().item())

Tổng p(x): 1.0 | Tổng q(x): 1.0

def speculative_Verify_samples(p, q, num_trials=200_000):
    # (1) Mô hình nháp đề xuất x ~ q(x)
    x_draft = torch.multinomial(q, num_samples=num_trials, replacement=True)

    # (2) Chấp nhận x với xác suất min(1, p(x) / q(x))
    accept_prob = torch.clamp(p[x_draft] / q[x_draft], max=1.0)
    u = torch.rand(num_trials)
    accepted = u < accept_prob

    # (3) Nếu bị bác bỏ, rút mẫu bù từ p'(x) = norm(max(0, p(x) - q(x)))
    residual = torch.clamp(p - q, min=0.0)
    p_res = residual / residual.sum()
    num_rej = int((~accepted).sum().item())
    x_final = x_draft.clone()
    x_final[~accepted] = torch.multinomial(
        p_res, num_samples=num_rej, replacement=True
    )
    return x_final, accepted.float().mean().item()


alpha_theory = torch.minimum(p_target, q_draft).sum().item()
samples, alpha_emp = speculative_Verify_samples(p_target, q_draft)
counts = torch.bincount(samples, minlength=len(tokens)).float()
p_emp = counts / len(samples)
max_err = (p_emp - p_target).abs().max().item()

print(f"Tỷ lệ chấp nhận lý thuyết α = Σ min(p, q): {alpha_theory*100:.2f}%")
print(f"Tỷ lệ chấp nhận thực nghiệm (200,000 mẫu):  {alpha_emp*100:.2f}%")
print(f"Sai lệch cực đại |p_spec(x) - p(x)|:       {max_err:.4f}\n")
print("Token  |  p_target(x) |  q_draft(x)  |  p_spec(x)")
print("-" * 49)
for i, tok in enumerate(tokens):
    print(
        f"{tok:6s} |     {p_target[i]:.4f}   |"
        f"    {q_draft[i]:.4f}    |    {p_emp[i]:.4f}"
    )

Tỷ lệ chấp nhận lý thuyết α = Σ min(p, q): 78.00%
Tỷ lệ chấp nhận thực nghiệm (200,000 mẫu):  77.96%
Sai lệch cực đại |p_spec(x) - p(x)|:       0.0014

Token  |  p_target(x) |  q_draft(x)  |  p_spec(x)
-------------------------------------------------
Hà     |     0.4200   |    0.2200    |    0.4191
Nội    |     0.2500   |    0.3800    |    0.2514
Sài    |     0.1500   |    0.1800    |    0.1498
Gòn    |     0.1000   |    0.0800    |    0.0999
Đà     |     0.0500   |    0.0900    |    0.0497
Nẵng   |     0.0300   |    0.0500    |    0.0301

Nếu muốn đào sâu hơn nữa vào kỹ thuật tối ưu hóa suy luận, bạn có thể thử sức với các bài tập mở rộng dưới đây.