Phục vụ nhiều người cùng lúc: batching và lượng tử hóa

Bài 3 kết luận rằng khi sinh chữ, mô hình gần như chỉ ngồi chờ bộ nhớ: mỗi token phải đọc toàn bộ 1,98 GB trọng số, còn bộ xử lý thì rảnh. Bài đó cũng để ngỏ một hệ quả: nếu đằng nào cũng phải đọc 1,98 GB, thì đọc một lần rồi dùng cho nhiều người cùng lúc. Bài này xây đúng điều đó: một "server" nhỏ trên mô hình tự viết ở bài 2 và 3, phục vụ nhiều người dùng bằng một bản trọng số duy nhất.

Lộ trình: batching và nó tốn bao nhiêu, vì sao xếp lịch kiểu "đợi cả nhóm xong" lãng phí và continuous batching sửa điều đó thế nào, KV cache quyết định số người dùng vừa bộ nhớ, rồi lượng tử hóa: làm trọng số nhỏ đi bằng cách làm tròn chúng, và đo xem làm tròn tới đâu thì mô hình hỏng.

Mọi con số trong bài đều đo trên máy thật, với Python 3.12, PyTorch 2.14, transformers 5.20, chạy bằng CPU với 2 luồng. Các thí nghiệm được giữ nhỏ (vài giây mỗi cái), nên con số tuyệt đối khác bài 3 một chút; điều đáng xem là tỉ lệ.

Batching: một lần đọc, nhiều người dùng

Ở bài 3, một bước decode đưa một token qua mô hình. Với nhiều người dùng, mỗi người có một token mới của riêng mình, ở một vị trí khác nhau, với một KV cache riêng. Mẹo nằm ở chỗ: trong mô hình, chỉ có attention là cần biết token đó thuộc về ai. Mọi phép nhân với trọng số (q, k, v, o, MLP) xử lý từng token độc lập, nên có thể gom token của mọi người thành một ma trận và nhân một lần:

def decode_step(next_ids: list[int], caches: list[dict]) -> torch.Tensor:
    """One decode step for many users at once: one new token each, each with its own cache and position."""
    x = W["embed_tokens.weight"][next_ids]                         # (users, 896)
    positions = [c[0][0].shape[1] for c in caches]                 # every user is at a different length
    for i in range(LAYERS):
        p = f"layers.{i}.self_attn."
        h = rms_norm(x, W[f"layers.{i}.input_layernorm.weight"])
        q = h @ W[p + "q_proj.weight"].T + W[p + "q_proj.bias"]    # one matrix multiply for all users:
        k = h @ W[p + "k_proj.weight"].T + W[p + "k_proj.bias"]    # the weights are read once
        v = h @ W[p + "v_proj.weight"].T + W[p + "v_proj.bias"]
        out = []
        for u, cache in enumerate(caches):                         # attention: each user looks at their own cache
            pos = torch.tensor([positions[u]])
            qu = rope(q[u].view(1, H, HD).transpose(0, 1), pos)
            ku = rope(k[u].view(1, KV, HD).transpose(0, 1), pos)
            vu = v[u].view(1, KV, HD).transpose(0, 1)
            ck, cv = torch.cat([cache[i][0], ku], 1), torch.cat([cache[i][1], vu], 1)
            cache[i] = (ck, cv)
            a = torch.nn.functional.scaled_dot_product_attention(qu, ck, cv, enable_gqa=True)
            out.append(a.transpose(0, 1).reshape(D))
        x = x + torch.stack(out) @ W[p + "o_proj.weight"].T
        x = x + mlp(rms_norm(x, W[f"layers.{i}.post_attention_layernorm.weight"]), i)
    return rms_norm(x, W["norm.weight"]) @ W["embed_tokens.weight"].T  # (users, 151936)

Các hàm rms_norm, rope, mlp là của bài 2, nguyên vẹn. Điều kiện để tin được: hai người dùng chạy chung một batch phải ra đúng từng token như khi mỗi người chạy riêng. Kiểm tra với hai câu hỏi khác nhau, 16 token mỗi người: giống hệt.

Thời gian một bước, khi càng nhiều người cùng chờ token kế tiếp:

Từ 2 lên 16 người, một bước chỉ chậm thêm khoảng 40%, trong khi số token sinh ra mỗi bước tăng 8 lần. Phục vụ 16 người cùng lúc cho gấp 6 lần số token mỗi giây so với phục vụ từng người một, với cùng một bản trọng số trong bộ nhớ. Riêng bước từ 1 lên 2 người thì thời gian nhảy gần gấp đôi rồi mới phẳng ra; có lẽ vì thư viện tính toán có một đường chạy riêng, nhanh hơn hẳn, cho trường hợp nhân với đúng một vector. Hình dạng chung thì đúng như bài 3 dự đoán: chi phí chính là đọc trọng số, và nó được chia đều cho mọi người trong batch.

Cái giá phải trả: mỗi người nhận chữ chậm hơn (36,5 xuống 13,7 token/s). Mọi hệ thống phục vụ LLM đều phải chọn một điểm trên đường cong này: batch lớn thì rẻ hơn cho mỗi token, batch nhỏ thì mỗi người thấy chữ chạy nhanh hơn.

Xếp lịch: đừng để cả nhóm chờ người chậm nhất

Batching đặt ra một câu hỏi mới: ai được vào batch, và khi nào? Cách đơn giản nhất là static batching: gom 4 yêu cầu thành một nhóm, chạy tới khi cả nhóm xong, rồi mới lấy nhóm tiếp theo. Vấn đề là câu trả lời dài ngắn khác nhau. Một câu "1 + 1 bằng mấy?" cần vài token, một câu chuyện cần hàng trăm. Trong một nhóm, người xong sớm vẫn chiếm chỗ cho tới khi người cuối cùng xong, và những yêu cầu đang xếp hàng phải đợi.

Continuous batching (cách mà vLLM và các server hiện đại làm) xếp lịch ở mức từng bước: ngay khi một người xong, chỗ của họ được trao cho yêu cầu kế tiếp trong hàng đợi, ở bước decode tiếp theo.

def continuous(requests, slots=4):
    queue, active = list(requests), []
    while queue or active:
        while queue and len(active) < slots:                      # a free slot: let the next request in now
            request = queue.pop(0)
            next_id, cache = prefill(request.prompt)
            active.append([request, next_id, cache, 1])
        out = decode_step([a[1] for a in active], [a[2] for a in active]).argmax(-1).tolist()
        for a, next_id in zip(active, out):
            a[1], a[3] = next_id, a[3] + 1
        for a in [a for a in active if a[3] >= a[0].max_tokens]:
            active.remove(a)                                      # finished: its slot frees up immediately

Đo với 12 yêu cầu đến cùng lúc, mỗi cái cần một số token khác nhau (8 tới 64, tổng cộng 381), tối đa 4 người một batch:

Cùng mô hình, cùng batch tối đa 4, chỉ khác cách xếp lịch: nhanh hơn 1,65 lần, và mỗi người chờ ít hơn một phần ba. Với static batching, nhóm có yêu cầu 64 token bắt ba người kia ngồi không trong gần hết thời gian; với continuous batching, không có slot nào để trống khi còn người đang xếp hàng.

Một chi tiết trong code ở trên: yêu cầu mới vào batch phải prefill prompt của nó trước. Prefill thì giới hạn bởi sức tính (bài 3), còn decode giới hạn bởi bộ nhớ, nên các server thật trộn chúng khéo hơn: cắt prefill của prompt dài thành từng mẩu và chạy xen với các bước decode của người khác, để một prompt dài không làm cả batch khựng lại.

KV cache quyết định bao nhiêu người vừa

Batching chia sẻ được trọng số, nhưng KV cache thì không chia sẻ được: mỗi người có cache riêng, lớn dần theo độ dài cuộc hội thoại. Bài 3 đã tính: 12 KiB mỗi token với Qwen2.5-0.5B, 56 KiB với bản 7B, 320 KiB với bản 72B. Bộ nhớ còn lại sau khi chứa trọng số chính là ngân sách cho cache, và nó quyết định batch lớn được bao nhiêu.

Một phép tính ví dụ (ước tính, không đo, bỏ qua bộ nhớ cho các phép tính tạm): Qwen2.5-7B ở bfloat16 trên một GPU 24 GB. Trọng số chiếm khoảng 15,2 GB, còn khoảng 8,8 GB cho cache. Mỗi cuộc hội thoại 4.096 token cần 56 KiB × 4.096 = 224 MiB, nên vừa khoảng 37 cuộc hội thoại cùng lúc. Nếu trọng số được nén xuống 4-bit (mục dưới), chúng chỉ còn khoảng 4,3 GB, cache có gần 20 GB, và số cuộc hội thoại tăng lên khoảng 84. Lượng tử hóa không chỉ làm mỗi bước nhanh hơn; nó còn mở chỗ cho batch lớn hơn.

Còn một kiểu lãng phí nữa: nếu dành sẵn cho mỗi người chỗ cho độ dài tối đa (32.768 token) trong khi phần lớn cuộc hội thoại chỉ dài vài trăm token, phần lớn bộ nhớ cache nằm trống. vLLM giải quyết bằng PagedAttention: chia cache thành các trang nhỏ cố định và cấp phát dần khi cuộc hội thoại dài ra, giống cách hệ điều hành quản lý bộ nhớ ảo.

Lượng tử hóa: làm tròn trọng số

Bài 3 cho thấy giảm trọng số từ 32-bit xuống 16-bit làm decode nhanh 1,7 lần, vì có ít byte hơn để đọc. Đi xa hơn nữa: thay mỗi số thực bằng một số nguyên nhỏ cộng một hệ số tỉ lệ dùng chung cho một nhóm.

def quantize(w: torch.Tensor, bits: int, group: int = 32):
    """Round each group of 32 weights to 2**bits levels; keep one float scale per group."""
    g = w.reshape(-1, group)
    qmax = 2 ** (bits - 1) - 1                                  # 127 for 8 bits, 7 for 4 bits
    scale = g.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / qmax
    q = (g / scale).round().clamp(-qmax - 1, qmax).to(torch.int8)  # the small integers that get stored
    return q, scale


def dequantize(q, scale, shape):
    return (q.float() * scale).reshape(shape)

Với mỗi nhóm 32 trọng số liên tiếp: tìm giá trị tuyệt đối lớn nhất, chia đều khoảng đó thành 2^bits mức, và làm tròn mỗi trọng số về mức gần nhất. Ở 8-bit có 256 mức; ở 4-bit chỉ còn 16; ở 2-bit chỉ còn 4 (−2, −1, 0, 1 nhân với hệ số). Thứ cần lưu là các số nguyên nhỏ, cộng một hệ số 16-bit cho mỗi 32 trọng số.

Để đo chất lượng, áp dụng quantize rồi dequantize cho mọi ma trận của Qwen2.5-0.5B trong mô hình tự viết, rồi đo hai thứ: số bit mỗi byte trên một đoạn 10 KB Wikipedia tiếng Việt (thước đo ở bài 4, càng thấp càng tốt), và câu trả lời tham lam cho "Thủ đô của Việt Nam là gì?":

• 8-bit: không mất gì đo được. Nhỏ hơn gần 4 lần, chất lượng giống hệt tới ba chữ số thập phân. 256 mức là thừa đủ cho trọng số của một mô hình.
• 4-bit: mất một ít. Nhỏ hơn 7 lần, số bit mỗi byte tăng 10%, câu trả lời vẫn đúng. Đây là mức mà phần lớn người chạy mô hình trên máy cá nhân dùng.
• 3-bit: sụp đổ. Số bit mỗi byte gấp đôi, ngang với xz ở bài 4, và mô hình chỉ in ra dấu sao.
• 2-bit: vô nghĩa. Tệ hơn cả gzip; mô hình nhả ra token ngẫu nhiên từ đủ thứ ngôn ngữ.

Cách làm tròn ở trên là cách đơn giản nhất (làm tròn tới mức gần nhất). Các phương pháp như GPTQ hay AWQ chọn cách làm tròn khôn hơn, dùng một ít dữ liệu mẫu để biết trọng số nào quan trọng, và giữ được chất lượng tốt hơn ở 4-bit và 3-bit. Mô hình lớn cũng thường chịu lượng tử hóa tốt hơn mô hình 0,5B ở đây. Nhưng hình dạng chung là như vậy: gần như miễn phí tới 8-bit, rẻ ở 4-bit, rồi tới một vách đá.

Còn tốc độ? Thí nghiệm trên làm tròn rồi lại đổi về số thực 32-bit trước khi nhân, nên nó chỉ đo chất lượng, không nhanh hơn. Để nhanh thật, phần mềm phải đọc thẳng các số nguyên nhỏ từ bộ nhớ và giải nén ngay trong lúc nhân (các thư viện như llama.cpp hay vLLM có sẵn những hàm như vậy). Khi đó, theo công thức ở bài 3 (token mỗi giây tối đa ≈ băng thông bộ nhớ ÷ byte trọng số), bản int8 0,53 GB có giới hạn trên gần gấp 4 lần bản float32, và bản int4 0,28 GB gấp khoảng 7 lần. Đây là ước tính từ công thức, không phải số đo.

Ghép lại: một server LLM thật làm gì

Mọi mảnh trong seri tới giờ đều xuất hiện trong một server LLM như vLLM:

Phía người dùng, server đó trông giống hệt chatbot ở seri Python: một endpoint nhận câu hỏi và stream từng token về. Khác biệt là phía sau endpoint, token của hàng chục người được sinh ra trong cùng một bước, bằng cùng một bản trọng số.

Tóm lại

Năm bài đã đi từ token tới server. Bài cuối cùng quay về phía người làm sản phẩm: dùng mô hình để làm một ứng dụng thật, với lời dặn hệ thống, gọi công cụ, vòng lặp agent, và trả lời câu hỏi dựa trên tài liệu của chính mình.