Inference cost at scale with napkin math：规模化推理成本的“餐巾纸数学”

背景

在将 AI 模型作为产品技术栈一部分提供服务时，开发者最常面临的挑战之一是确定 GPU 集群的规模上限。为了回答“我的集群能支撑多少用户”以及“每个用户的推理成本是多少”这类问题，通常需要进行复杂的工程估算。

这篇文章通过“餐巾纸数学”（Napkin Math，指基于基本假设的快速粗略估算）的方式，从硬件规格和模型架构的基础知识出发，逐步推导出单个用户的美元成本。文章旨在帮助那些熟悉 GPU 和大型语言模型（LLM）的工程师，通过简化的数学模型，直观地理解推理过程中的计算瓶颈、内存带宽限制以及 KV-Cache 优化的巨大价值。

核心内容

文章通过一系列逻辑严密的步骤，从基础的矩阵乘法成本推导至最终的单用户成本估算。

1. 单个 GPU 的资源指标

在任何 GPU 的规格说明书中，主要关注两个核心指标：

• 峰值吞吐量（Peak Throughput）：每秒浮点运算次数，通常以 TFLOP/s 为单位（$1 \text{ TFLOP/s} = 10^{12} \text{ ops/sec}$）。
• 内存带宽（Memory Bandwidth）：数据从全局内存（VRAM）移动到寄存器（SRAM）的速度，通常以 TB/sec 为单位。

为了简化计算，文章假设使用 FP-8 量化（也可轻松调整为 FP-16）。

2. 矩阵乘法（Matmul）的成本

LLM 的核心操作是大规模矩阵乘法。假设两个矩阵 $A_{N \times d}$ 和 $B_{d \times M}$ 相乘得到 $O_{N \times M}$。

• 基础成本：对于 $O$ 中的每个元素，需要加载 $A$ 和 $B$ 的元素，相乘后累加。计算 $(N,d) \times (d,M)$ 的矩阵乘积，需要 $2NMd$ 次内存访问和 $2NMd$ 次浮点运算。
• 分块优化（Tiling）：通过分块技术，内存访问次数可降低至约 $d(N+M)$。

3. 语言模型概览

LLM 本质上是自回归模型：接收 $N$ 个词的序列，生成第 $N+1$ 个词。

• 每个词表示为 $d$ 维向量。
• 通过反复应用“注意力机制”（Attention），预测下一个词。
• 前向传播过程包括：输入嵌入、多层注意力计算、最终层输出转换为词概率分布（通过 $W_{vocab}$ 矩阵和 Softmax），并采样得到下一个 token。

4. 注意力机制详解

注意力机制的计算公式为： $$ Attention(Q, K, V) = \text{softmax}(Q K^T / \sqrt{d}) V $$ 其中 $Q, K, V$ 分别由输入 $X$ 与权重矩阵 $W_Q, W_K, W_V$ 相乘得到。

在实际推理中，为了并行处理 $B$ 个并发对话，输入变为 $X \in \mathbb{R}_{B \times N \times d}$。

5. 内存带宽瓶颈：未优化前的困境

文章指出，当前的实现存在严重的内存读取问题。以生成单个 token 为例：

• 假设上下文长度 $N \approx 200,000$，维度 $d = 8,192$。
• 对于单次 $K @ W_k$ 矩阵乘法，需要 $26$ 万亿次浮点运算和 $17$ 亿次内存访问。
• 关键洞察：计算量远大于内存访问量（计算量比内存访问多四个数量级）。这意味着 GPU 需要等待数万个周期才能完成当前批次的计算，内存带宽并非主要瓶颈，但计算效率极低，因为每次生成新 token 时，都要重新处理整个历史对话。

6. KV-Cache 优化：降低计算成本

为了避免重复计算已处理的 token，推理引擎（如 vLLM）会在每一层缓存中间输出 $K$ 和 $V$，存储在 VRAM 中的 KV Cache 区域。

• 优化效果：有了 KV-Cache，每次前向传播只需处理最新生成的一个 token（$X \in \mathbb{R}_{1 \times d}$），而不是整个历史序列。
• 新的成本计算：
  • 对于批次大小 $B$，内存访问次数降至约 $2620$ 万次。
  • 浮点运算降至约 $5240$ 万次。
  • 效率提升：现在每进行一次内存访问，只需执行两次操作（而非之前的上万次）。这是一个巨大的效率飞跃。

7. 单 Token 成本与 GPU 规格

文章以 NVIDIA B200 为例进行估算：

• 内存带宽：$8 \text{ TB/s}$ ($8 \times 10^{12}$ 字节/秒)。
• 计算强度：$4500 \text{ TFLOP/s}$。

通过结合上述优化后的计算量与 GPU 规格，可以进一步推导出每秒生成的 Token 数（Tokens Per Second）以及每个用户的美元成本。

关键要点

• 计算与内存的权衡：在未使用 KV-Cache 的情况下，LLM 推理面临巨大的冗余计算问题，因为每次生成新 token 都要重新计算整个上下文的注意力机制，导致计算资源浪费严重。
• KV-Cache 的核心价值：通过缓存 Key 和 Value 矩阵，推理引擎将每次前向传播的计算复杂度从 $O(N)$ 降低到 $O(1)$（仅针对最新 token），从而极大地提高了吞吐量并降低了延迟。
• 自回归特性：LLM 的自回归性质决定了其生成过程是串行的，但通过批处理（Batching）和 KV-Cache 技术，可以在单张 GPU 上并行服务多个用户。
• 硬件规格的重要性：GPU 的内存带宽和峰值吞吐量是决定推理成本的关键硬件指标。NVIDIA B200 等高端 GPU 的高带宽和高算力使得大规模低成本推理成为可能。
• 餐巾纸数学的实用性：通过简化的数学模型（如假设 FP-8 量化、忽略 RoPE 等细节），可以快速估算出系统的理论上限，为架构决策提供直观依据。

意义与影响

这篇文章为 AI 基础设施工程师和产品经理提供了一个清晰的框架，用于理解 LLM 推理背后的经济和技术逻辑。

1. 成本透明化：通过“餐巾纸数学”，开发者可以不再依赖黑盒式的云服务报价，而是基于硬件规格自行估算推理成本，从而更精准地进行产品定价和资源规划。
2. 优化方向明确：文章强调了 KV-Cache 在降低计算冗余方面的决定性作用，这解释了为什么 vLLM 等现代推理引擎成为行业标准。它指导工程师将优化重点放在内存管理和缓存策略上，而非仅仅关注模型架构本身。
3. 规模化决策支持：帮助团队回答“我的 GPU 集群能支撑多少用户”这一核心业务问题。通过量化每个 token 的计算和内存需求，可以更科学地评估硬件投资回报率（ROI）。
4. 技术普及：将复杂的矩阵代数和注意力机制简化为易于理解的数学步骤，降低了理解 LLM 推理性能瓶颈的门槛，促进了 AI 工程领域的知识共享。

总之，这篇解读不仅是一篇技术教程，更是一份关于如何以最低成本实现大规模 AI 服务落地的实用指南。它揭示了在追求模型智能的同时，工程优化对于商业可行性的关键作用。