AI Inference Engineering: giải thích chi tiết kèm ví dụ

1. Tổng quan

Khi bạn gọi API của một chatbot và chờ câu trả lời chạy dần ra màn hình, bạn đang chứng kiến quá trình inference của một Large Language Model (LLM). Bài này giải thích inference khác gì với training, vì sao nó tốn tài nguyên khủng khiếp, và những kỹ thuật kỹ sư dùng để làm nó nhanh hơn, rẻ hơn: KV cache, continuous batching, quantization, speculative decoding. Đọc xong, bạn sẽ hiểu được các con số TTFT/TPOT trong benchmark, và biết viết code Java đo hiệu năng một inference endpoint thực tế.

Lưu ý: Bài không đi sâu vào nội bộ transformer (attention, feed-forward…). Trọng tâm là góc nhìn kỹ thuật hạ tầng — cách một request “chạy” qua hệ thống inference.

2. Inference là gì?

Training là quá trình học trọng số (weights) từ dữ liệu — tốn hàng nghìn GPU trong nhiều tuần, nhưng chỉ làm một lần. Inference là lúc model đã học xong, bạn đưa prompt vào và nhận output — việc này xảy ra mỗi lần có request, với hàng triệu request mỗi ngày trên production. Vì vậy, dù training “đắt” hơn về tổng chi phí một lần, inference mới là thứ quyết định chi phí vận hành dài hạn và trải nghiệm người dùng (latency).

Điểm khác biệt quan trọng: trong training, model nhìn thấy toàn bộ câu để tính loss một lượt. Trong inference sinh văn bản (text generation), model phải sinh từng token một, và token sau phụ thuộc vào token trước — gọi là autoregressive decoding.

2.1. Hai giai đoạn: Prefill và Decode

Một request inference luôn trải qua 2 giai đoạn:

• Prefill: model đọc toàn bộ prompt đầu vào cùng lúc, tính attention cho tất cả token input. Giai đoạn này tận dụng tốt GPU vì xử lý song song (compute-bound).
• Decode: model sinh token output lần lượt, mỗi lần một token, dựa vào toàn bộ context trước đó. Giai đoạn này chạy tuần tự, mỗi bước chỉ tính cho 1 token mới (memory-bound — nghẽn ở băng thông đọc bộ nhớ, không phải ở phép tính).

Minh họa bằng Java vòng lặp decode đơn giản (giả lập, không phải model thật) để thấy rõ bản chất tuần tự:

public class AutoregressiveDecoder {

    public List<String> generate(List<String> prompt, int maxNewTokens) {
        List<String> context = new ArrayList<>(prompt);
        List<String> generated = new ArrayList<>();

        for (int step = 0; step < maxNewTokens; step++) {
            String nextToken = predictNextToken(context); // 1 forward pass
            if (nextToken.equals("<EOS>")) break;

            context.add(nextToken);   // token mới trở thành input của bước sau
            generated.add(nextToken);
        }
        return generated;
    }

    private String predictNextToken(List<String> context) {
        // giả lập: model thật chạy forward pass qua toàn bộ transformer
        return "token_" + context.size();
    }
}

Mỗi lần gọi predictNextToken, model phải đọc lại toàn bộ context đã có để tính attention — nếu không tối ưu, chi phí tăng dần theo độ dài câu. Đây chính là động lực ra đời của KV cache ở phần 4.

3. Đo hiệu năng inference: TTFT và TPOT

Hai chỉ số quan trọng nhất khi benchmark một LLM serving:

• TTFT (Time To First Token): thời gian từ lúc gửi request đến khi nhận token đầu tiên — phản ánh độ trễ của giai đoạn prefill.
• TPOT (Time Per Output Token): thời gian trung bình để sinh mỗi token tiếp theo trong giai đoạn decode.

public record InferenceMetrics(Duration ttft, Duration avgTpot, int totalTokens) {

    public Duration totalLatency() {
        return ttft.plus(avgTpot.multipliedBy(totalTokens - 1));
    }
}

Hàm totalLatency() cho thấy công thức thực tế: tổng thời gian ≈ TTFT (1 lần) + TPOT × (số token còn lại). Với một câu trả lời 200 token, TPOT chỉ cần chậm 10ms/token cũng kéo thêm 2 giây tổng thời gian — đây là lý do TPOT bị soi kỹ hơn TTFT trong các benchmark serving.

Ví dụ số liệu thực tế (model 7B, batch size 1, GPU A100):
TTFT  ≈ 80ms
TPOT  ≈ 15ms/token
200 token → 80ms + 15ms × 199 ≈ 3.07s

4. KV Cache: tránh tính lại những gì đã biết

Nếu không cache, mỗi bước decode phải tính lại attention cho toàn bộ chuỗi token từ đầu — độ phức tạp tăng theo bình phương độ dài. KV cache lưu lại key/value vector của các token đã xử lý, để bước sau chỉ cần tính cho token mới, rồi tái sử dụng cache cũ.

public class KvCache {

    private final Map<Integer, float[]> keyCache = new HashMap<>();
    private final Map<Integer, float[]> valueCache = new HashMap<>();

    public void store(int position, float[] key, float[] value) {
        keyCache.put(position, key);
        valueCache.put(position, value);
    }

    public List<float[]> allKeysUpTo(int position) {
        return IntStream.rangeClosed(0, position)
            .mapToObj(keyCache::get)
            .toList();
    }
}

store được gọi một lần cho mỗi token mới sinh ra; allKeysUpTo cho phép bước decode kế tiếp lấy lại toàn bộ key đã tính mà không cần forward pass lại từ đầu. Đây là lý do một request dài (context dài) ngốn VRAM rất nhanh — cache này tỉ lệ thuận với batch_size × sequence_length × num_layers × hidden_size.

Lưu ý: KV cache là nguyên nhân chính khiến context window dài (32K, 128K token) tốn RAM GPU khủng khiếp, chứ không phải compute. Đây cũng là lý do các kỹ thuật như paged attention (dùng trong vLLM) ra đời — quản lý cache như bộ nhớ ảo (page) để tránh phân mảnh.

5. Continuous batching: gom request để tăng throughput

GPU xử lý theo batch sẽ hiệu quả hơn xử lý từng request đơn lẻ. Nhưng mỗi request có độ dài sinh khác nhau — request A xong ở token 50, request B còn tiếp tục đến token 300. Static batching (chờ cả batch xong mới nhận batch mới) sẽ lãng phí GPU vì phải đợi request chậm nhất. Continuous batching giải quyết bằng cách: ngay khi một request trong batch hoàn thành, “ghế” của nó được thay bằng request mới trong hàng đợi, không cần chờ cả batch.

public class ContinuousBatchScheduler {

    private final int batchSlots;
    private final Queue<InferenceRequest> waiting = new LinkedList<>();
    private final List<InferenceRequest> active;

    public ContinuousBatchScheduler(int batchSlots) {
        this.batchSlots = batchSlots;
        this.active = new ArrayList<>(Collections.nCopies(batchSlots, null));
    }

    public void onStep() {
        for (int slot = 0; slot < batchSlots; slot++) {
            InferenceRequest req = active.get(slot);
            if (req == null || req.isFinished()) {
                active.set(slot, waiting.poll()); // lấp ngay request mới
            } else {
                req.decodeNextToken();
            }
        }
    }
}

onStep() chạy mỗi bước decode: nếu slot nào trống hoặc request vừa xong, nó lập tức nhận request chờ kế tiếp từ waiting. Nhờ vậy, GPU luôn chạy full batch, không có slot “rỗi” chờ request chậm — đây là cơ chế lõi trong vLLM, TGI (Text Generation Inference) và TensorRT-LLM.

6. Quantization: đánh đổi độ chính xác lấy tốc độ

Model gốc thường lưu trọng số ở FP16 hoặc BF16 (16-bit). Quantization giảm số bit lưu trữ (INT8, INT4) để giảm dung lượng VRAM và tăng throughput, đổi lại độ chính xác giảm nhẹ.

Model 7B tham số, so sánh dung lượng:
- FP16: ~14 GB VRAM
- INT8: ~7 GB VRAM
- INT4: ~3.5 GB VRAM

Khi serve bằng các engine như vLLM, bạn chỉ cần khai báo cấu hình:

model: meta-llama/Llama-3-8B
quantization: awq
dtype: float16
gpu_memory_utilization: 0.9

quantization: awq chỉ định dùng phương pháp AWQ (Activation-aware Weight Quantization) — một kỹ thuật INT4 giữ lại độ chính xác tốt hơn so với quantize “ngây thơ” bằng cách ưu tiên giữ nguyên các trọng số quan trọng nhất.

Lưu ý: Quantization quá sâu (INT4 trở xuống) có thể làm model “ảo giác” (hallucination) nhiều hơn hoặc giảm khả năng reasoning. Luôn benchmark lại chất lượng output (không chỉ tốc độ) sau khi quantize, đặc biệt với các task cần độ chính xác cao như code generation hay tính toán.

7. Speculative decoding: đoán trước để tiết kiệm bước

Decode tuần tự chậm vì mỗi bước chỉ sinh 1 token. Speculative decoding dùng một model nhỏ, nhanh (“draft model”) để đoán trước vài token tiếp theo, rồi model lớn (“target model”) chỉ cần xác nhận (verify) các token đó trong một lần forward pass thay vì nhiều lần.

Draft model đoán: "Hà" "Nội" "là" "thủ" "đô"
Target model verify cả 5 token cùng lúc trong 1 forward pass:
  - Nếu 4/5 token khớp với phân phối của model lớn → giữ 4 token, bỏ token thứ 5, decode lại từ đó.
  - Nếu sai từ đầu → chỉ tốn đúng 1 forward pass như decode thường, không mất gì.

Kỹ thuật này không làm giảm chất lượng output (vì luôn verify bằng target model), chỉ giảm số lần phải chạy forward pass tuần tự trên model lớn — hiệu quả nhất khi draft model đoán đúng nhiều.

8. Thực hành: đo TTFT/TPOT từ một LLM serving endpoint

Giả sử bạn có một vLLM server chạy OpenAI-compatible API ở http://localhost:8000. Thêm dependency cho WebClient:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-webflux</artifactId>
</dependency>

Viết client đo TTFT bằng streaming response:

@Service
public class LlmBenchmarkClient {

    private final WebClient webClient;

    public LlmBenchmarkClient(WebClient.Builder builder) {
        this.webClient = builder.baseUrl("http://localhost:8000").build();
    }

    public InferenceMetrics measure(String prompt) {
        Instant start = Instant.now();
        AtomicReference<Instant> firstTokenTime = new AtomicReference<>();
        AtomicInteger tokenCount = new AtomicInteger();

        webClient.post()
            .uri("/v1/completions")
            .bodyValue(Map.of("prompt", prompt, "stream", true, "max_tokens", 100))
            .retrieve()
            .bodyToFlux(String.class)
            .doOnNext(chunk -> {
                if (firstTokenTime.get() == null) {
                    firstTokenTime.set(Instant.now());
                }
                tokenCount.incrementAndGet();
            })
            .blockLast();

        Duration ttft = Duration.between(start, firstTokenTime.get());
        Duration total = Duration.between(start, Instant.now());
        Duration avgTpot = tokenCount.get() > 1
            ? total.minus(ttft).dividedBy(tokenCount.get() - 1)
            : Duration.ZERO;

        return new InferenceMetrics(ttft, avgTpot, tokenCount.get());
    }
}

doOnNext được gọi mỗi khi một chunk (token) mới về — lần gọi đầu tiên ghi lại thời điểm firstTokenTime chính là TTFT. Sau khi stream kết thúc (blockLast), ta tính TPOT trung bình bằng tổng thời gian trừ TTFT, chia cho số token còn lại.

8.1. Test với JUnit 5 và AssertJ

@SpringBootTest
class LlmBenchmarkClientTest {

    @Autowired
    LlmBenchmarkClient client;

    @Test
    void shouldMeetLatencySla() {
        InferenceMetrics metrics = client.measure("Giải thích KV cache trong 2 câu.");

        assertThat(metrics.ttft()).isLessThan(Duration.ofMillis(500));
        assertThat(metrics.avgTpot()).isLessThan(Duration.ofMillis(50));
        assertThat(metrics.totalTokens()).isGreaterThan(0);
    }
}

Test này không kiểm tra nội dung câu trả lời, chỉ kiểm tra SLA về độ trễ — đúng trọng tâm của “inference engineering”: đảm bảo hệ thống phản hồi đủ nhanh, không phải đánh giá chất lượng model (việc đó thuộc về đánh giá chất lượng, khác phạm vi bài này).

Lưu ý: Test này cần một server LLM thật đang chạy (hoặc mock WebClient bằng MockWebServer). Trong CI, nên tách riêng test loại này thành integration test, không chạy cùng unit test thông thường vì phụ thuộc hạ tầng bên ngoài.

9. Lỗi thường gặp và khi nào không nên tự xây inference engine

• Tự viết lại batching/KV cache từ đầu: trừ khi bạn làm nghiên cứu hạ tầng ML chuyên sâu, đừng tự build lại những gì vLLM, TGI, TensorRT-LLM đã tối ưu kỹ — bạn sẽ mất hàng tháng chỉ để đạt lại hiệu năng họ có sẵn.
• Benchmark chỉ bằng throughput, quên latency đuôi (tail latency): throughput cao nhưng p99 latency tệ vẫn làm trải nghiệm người dùng xấu. Luôn đo p50/p95/p99, không chỉ trung bình.
• Quantize mà không test lại chất lượng: như đã nói ở phần 6, giảm bit luôn có đánh đổi — đo tốc độ mà bỏ qua đo chất lượng là sai lầm phổ biến.
• Dùng static batching cho traffic thực tế biến động: nếu request có độ dài output chênh lệch lớn, static batching sẽ lãng phí GPU nghiêm trọng — nên chọn engine hỗ trợ continuous batching ngay từ đầu.
• Không tách biệt prefill-heavy và decode-heavy workload: một số hệ thống production tách node xử lý prefill và node xử lý decode riêng (disaggregated serving) để tối ưu tài nguyên — nếu traffic của bạn có pattern rõ (ví dụ RAG với prompt rất dài, output ngắn), đáng xem xét kiến trúc này.

10. Kết luận

Inference engineering là bài toán tối ưu giữa latency, throughput và chi phí GPU — không phải chuyện huấn luyện model. Nắm được prefill/decode, KV cache, continuous batching, quantization và speculative decoding giúp bạn đọc hiểu benchmark, chọn đúng serving engine (vLLM, TGI…), và đo lường hệ thống của mình bằng những chỉ số đúng (TTFT, TPOT) thay vì chỉ nhìn “nó chạy nhanh hay chậm”. Bước tiếp theo tự nhiên là thử đổi quantization hoặc gpu_memory_utilization trong config thật và quan sát các chỉ số này thay đổi ra sao trên workload của riêng bạn.