AI 推理工程：详细解释与示例

1. 概览

当你调用某个聊天机器人的 API 并等待回复逐个显示在屏幕上时，你正在见证 Large Language Model（LLM）的 推理 过程。本文解释推理与训练的区别、为什么它会消耗大量资源，以及工程师用来让它更快、更便宜的技术：KV cache、continuous batching、quantization、speculative decoding。读完这篇文章后，你将理解基准测试中的 TTFT/TPOT 数字，以及如何编写 Java 代码来测量实际推理端点的性能。

注意： 本文不深入讨论 transformer 的内部（attention、feed-forward 等）。重点是基础设施视角——一个请求如何通过推理系统“运行”。

2. 推理是什么？

训练是从数据中学习权重（weights）的过程——需要数千个 GPU 花费数周时间，但只做一次。推理是模型已学完后，你输入 prompt 并获得输出的过程——这在每次有请求时发生，生产环境中每天有数百万个请求。因此，虽然训练在一次性成本上更“昂贵”，但推理才是决定长期运营成本和用户体验（延迟）的因素。

重要的区别在于：在训练中，模型一次看到整个句子来计算损失。在推理文本生成中，模型必须逐个生成 token，后面的 token 依赖于前面的 token——这被称为 autoregressive decoding。

2.1. 两个阶段：Prefill 和 Decode

一个推理请求总是经历 2 个阶段：

• Prefill：模型一次性读取整个输入 prompt，为所有输入 token 计算注意力。这个阶段充分利用 GPU，因为处理是并行的（compute-bound）。
• Decode：模型逐个生成输出 token，每次一个，基于之前的所有上下文。这个阶段按顺序运行，每一步只计算 1 个新 token（memory-bound——受内存带宽限制，而不是计算能力）。

用简单的 Java 循环来说明 decode 的顺序特性（模拟，不是真实模型）：

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 次前向传递
            if (nextToken.equals("<EOS>")) break;

            context.add(nextToken);   // 新 token 成为下一步的输入
            generated.add(nextToken);
        }
        return generated;
    }

    private String predictNextToken(List<String> context) {
        // 模拟：真实模型运行前向传递遍历整个 transformer
        return "token_" + context.size();
    }
}

每次调用 predictNextToken 时，模型都必须重新读取所有现有的上下文来计算注意力——如果不优化，成本会随着句子长度而增加。这正是第 4 部分中 KV cache 出现的动力。

3. 测量推理性能：TTFT 和 TPOT

对 LLM serving 进行基准测试时的两个最重要指标：

• TTFT（Time To First Token）：从发送请求到收到第一个 token 的时间——反映 prefill 阶段的延迟。
• TPOT（Time Per Output Token）：在 decode 阶段生成每个后续 token 的平均时间。

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

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

totalLatency() 函数展示了实际公式：总时间 ≈ TTFT（1 次）+ TPOT × （剩余 token 数）。对于 200 个 token 的响应，TPOT 如果只是 10ms/token 慢也会增加 2 秒的总时间——这就是为什么 TPOT 在 serving 基准测试中比 TTFT 受到更多关注。

实际数据示例（7B 模型，batch size 1，GPU A100）：
TTFT  ≈ 80ms
TPOT  ≈ 15ms/token
200 token → 80ms + 15ms × 199 ≈ 3.07s

4. KV Cache：避免重新计算已知的内容

如果不缓存，每个 decode 步骤都必须为整个 token 序列重新计算注意力——复杂度按长度的平方增加。KV cache 保存已处理 token 的 key/value 向量，以便下一步只需为新 token 计算，然后重用旧缓存。

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 为每个新生成的 token 调用一次；allKeysUpTo 允许下一个 decode 步骤检索所有已计算的 key，无需从头再做前向传递。这就是为什么长请求（长上下文）会快速消耗 VRAM——此缓存与 batch_size × sequence_length × num_layers × hidden_size 成正比。

注意： KV cache 是导致长上下文窗口（32K、128K token）消耗大量 GPU RAM 的主要原因，而不是计算。这也是为什么 paged attention 等技术（在 vLLM 中使用）出现的原因——像虚拟内存（分页）一样管理缓存以避免碎片。

5. Continuous batching：聚合请求以提高吞吐量

GPU 按批处理会比逐个处理请求更高效。但每个请求的生成长度不同——请求 A 在 50 个 token 处完成，请求 B 继续到 300 个 token。Static batching（等待整个批处理完成才接收新批）会浪费 GPU，因为必须等最慢的请求。Continuous batching 解决方案是：当批中的一个请求完成时，它的“槽位”立即被队列中的新请求替换，无需等待整个批。

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()); // 立即填充新请求
            } else {
                req.decodeNextToken();
            }
        }
    }
}

onStep() 在每个 decode 步骤运行：如果任何槽位为空或请求刚完成，它立即从 waiting 获取下一个等待请求。这样，GPU 始终运行完整批，没有“空闲”槽位等待慢请求——这是 vLLM、TGI（Text Generation Inference）和 TensorRT-LLM 的核心机制。

6. Quantization：牺牲精度换取速度

原始模型通常以 FP16 或 BF16（16 位）存储权重。Quantization 减少存储位数（INT8、INT4）来减少 VRAM 占用并增加吞吐量，代价是精度略微下降。

7B 参数模型，容量比较：
- FP16: ~14 GB VRAM
- INT8: ~7 GB VRAM
- INT4: ~3.5 GB VRAM

使用 vLLM 等引擎进行服务时，只需声明配置：

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

quantization: awq 指定使用 AWQ（Activation-aware Weight Quantization）方法——一种 INT4 技术，通过优先保留最重要的权重，保持比“朴素”量化更好的精度。

注意： 过度量化（INT4 及以下）可能导致模型产生更多“幻觉”（hallucination）或推理能力下降。量化后始终要重新基准测试输出质量（不仅仅是速度），特别是对于需要高精度的任务，如代码生成或数学计算。

7. Speculative decoding：预测以节省步骤

顺序解码很慢，因为每一步只生成 1 个 token。Speculative decoding 使用一个小的、快速的“draft model”来预测接下来的几个 token，然后大模型“target model”只需在一次前向传递中验证这些 token，而不是多次。

Draft model 预测："河" "内" "是" "首" "都"
Target model 在 1 次前向传递中同时验证全部 5 个 token：
  - 如果 4/5 token 与大模型的分布匹配 → 保留 4 个 token，丢弃第 5 个，从那里开始重新解码。
  - 如果从一开始就错了 → 只消耗 1 次前向传递，和普通解码一样，没有浪费。

这种技术不会降低输出质量（因为始终用 target model 验证），只是减少必须在大模型上顺序运行的前向传递次数——当 draft model 预测准确时效果最好。

8. 实践：从 LLM serving 端点测量 TTFT/TPOT

假设你有一个 vLLM 服务器在 http://localhost:8000 运行 OpenAI 兼容的 API。添加 WebClient 依赖：

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

编写一个用流式响应测量 TTFT 的客户端：

@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());
    }
}

每当一个新的 chunk（token）到达时，doOnNext 被调用——第一次调用记录 firstTokenTime 时刻就是 TTFT。在流结束后（blockLast），我们通过总时间减去 TTFT，再除以剩余 token 数来计算平均 TPOT。

8.1. 使用 JUnit 5 和 AssertJ 进行测试

@SpringBootTest
class LlmBenchmarkClientTest {

    @Autowired
    LlmBenchmarkClient client;

    @Test
    void shouldMeetLatencySla() {
        InferenceMetrics metrics = client.measure("在 2 句话内解释 KV cache。");

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

这个测试不检查回复的内容，只检查延迟 SLA——这正是“推理工程”的重点：确保系统响应足够快，而不是评估模型质量（那是质量评估，超出了本文的范围）。

注意： 这个测试需要一个实际运行的 LLM 服务器（或用 MockWebServer mock WebClient）。在 CI 中，应该将这类测试分离为集成测试，不要与常规单元测试一起运行，因为它依赖外部基础设施。

9. 常见错误和何时不应自建推理引擎

• 从头重写 batching/KV cache：除非你在做深度 ML 基础设施研究，否则不要尝试重新构建 vLLM、TGI、TensorRT-LLM 已经精心优化的东西——你需要花数月时间才能达到他们现有的性能。

• 仅用吞吐量进行基准测试，忽视尾部延迟（tail latency）：高吞吐量但 p99 延迟差仍会导致用户体验糟糕。始终测量 p50/p95/p99，而不仅仅是平均值。

• 量化后不重新测试质量：如第 6 部分所述，减少位数总是有代价的——测量速度而忽视质量测量是常见错误。

• 对于实际变化的流量使用 static batching：如果请求的输出长度差异很大，static batching 会严重浪费 GPU——应该从一开始就选择支持 continuous batching 的引擎。

• 不区分 prefill-heavy 和 decode-heavy 工作负载：一些生产系统分离处理 prefill 和 decode 的节点（分离服务），以优化资源——如果你的流量有明确的模式（例如 RAG 提示很长，输出很短），值得考虑这种架构。

10. 结论

推理工程是在延迟、吞吐量和 GPU 成本之间的优化问题——不是关于训练模型的。掌握 prefill/decode、KV cache、continuous batching、quantization 和 speculative decoding 将帮助你理解基准测试，选择正确的服务引擎（vLLM、TGI 等），以及用正确的指标（TTFT、TPOT）测量你的系统，而不仅仅说“它快还是慢”。下一步自然是尝试在真实配置中改变 quantization 或 gpu_memory_utilization，观察这些指标在你自己的工作负载上如何变化。