如果您的产品堆栈中集成了 AI 模型，您可能曾好奇过，您的 GPU 集群究竟能支撑到什么规模。通过对您的硬件和模型架构进行一些基础的了解，我们就可以在“餐巾纸”上快速估算出每位用户的美元成本。
如果您对 GPU 和/或 LLM（大型语言模型）的原理比较熟悉，请使用以下索引跳转到相关的章节：
单张 GPU 的资源配置 矩阵乘法的成本 语言模型概述 注意力机制的深度解析 通过 KV-Cache 减少计算量 一个 token 的成本是多少？ 现实中能为多少用户提供服务？ 在单张 GPU 上为数百用户进行优化 Tokens Per Second (每秒 Token 数) 每用户美元成本
### 单张 GPU 的资源配置 (Resources on a single GPU)
在任何一张 GPU 的规格表（spec-sheet）上，您都可以找到以下指标：
峰值吞吐量 (Peak throughput)： 指每秒浮点运算次数（Floating-point operations per second）。通常以 TeraFLOPs (TFLOP/s) 表示（1 TFLOP/s = $10^{12}$ 次运算/秒）。
内存带宽 (Memory bandwidth)： 指数据从全局内存（VRAM）传输到寄存器（SRAM）所能移动的数据量。通常以 TB/秒 (TB/sec) 表示。
在计算吞吐量时，我们假设采用 FP-8 量化（FP-8 quantization），但如果您想调整计算以适应 FP-16，过程也非常简单。
### 矩阵乘法的成本 (Cost of a Matrix Multiplication)
如果您点开了这篇文章，您应该知道 AI 模型在海量的矩阵上执行了大量的矩阵乘法（matrix multiplications）。因此，我们从计算一次矩阵乘法的成本入手，这理应是意料之中的。
假设我们有两个矩阵：$A_{N imes d}$ 和 $B_{d imes M}$。 设它们的乘积为矩阵 $O_{N imes M}$。
根据高中代数知识，我们知道 $O$ 中的每个元素 $O^{i,k}$ 可以通过以下方式计算： $$O^{i,k} = sum_{j=1}^{d} A^{i,j} B^{j,k}$$
在这个过程中，我们获得了关于矩阵乘法“成本”的第一个洞察。对于每一个 $O^{i,k}$，我们需要从初始值 0 开始，执行以下操作：
1. 从内存中加载 $A^{i,j}$。 2. 从内存中加载 $B^{j,k}$。 3. 将两者相乘。 4. 将第 3 步的结果加到累积和上。
这个过程总共需要对每个元素重复 $d$ 次。
因此，一次 $(N,d) imes (d,M)$ 矩阵乘法的成本是： $2NMd$ 次 内存访问 (memory accesses) $2NMd$ 次 浮点运算 (floating-point operations)
通过一种称为“分块”（tiling）的优化技术，内存访问量可以减少到大约 $d(N+M)$。虽然这里不需要深入细节，但对于好奇的读者，Alvin 的博客文章中有详细的解释。
### 语言模型概述 (An Overview of Language Models)
从本质上讲，LLM 是一个简单的系统——它接收一个词元序列 (sequence of tokens)，并生成下一个词元。每个词元都被表示为一个 $d$ 维的向量。
通过重复应用一个称为“注意力机制”（attention）的函数（稍后会解释），模型来预测下一个词元。一次前向传播（forward pass）大致如下所示：
输入 (input) $# y = A (N imes d)$ 矩阵 $#$ 每个层 (each layer) 的网络： $ ext{attention}(y)$ $#$ 将最后一层的输出转换为词元概率 (Convert the final layer's output to word-probs)。 $# W_{ ext{vocab}}$ 是一个大小为 $d imes ext{vocab\_len}$ 的矩阵，其中 $ ext{vocab\_len}$ 是模型词汇表中所有词元的数量。 $ ext{token\_probs} = ext{softmax}(y W_{ ext{vocab}})$ $ ext{next\_tok} = ext{token\_probs}( ext{argmax}( ext{token\_probs}))$ $# ext{next\_tok}$ 是一个 $(1 imes d)$ 向量
这也是为什么 LLM 被称为自回归 (auto-regressive) 的原因。它可以持续对自身的输出进行多次前向传播，直到生成一个 $ ext{
这是一个简化的概述，其中省略了 RoPE（旋转位置编码）、中间的 MLP（多层感知机）层、末端的词元采样（token sampling）以及更多细节。
作为一个有趣的练习：您应该尝试将这些因素计入成本，并通过费米估算 (Fermi estimation) 来验证我们的数学模型是否依然成立。
### 注意力机制的深度解析 (Attention in Greater Detail)
我假设您对注意力机制有一定的了解，这里只提供一个回顾。
如前所述，输入是一个矩阵 $X in mathbb{R}^{N imes d}$，其中 $X_i$ 是一个 $d$ 维的单个向量。
对于网络中的每一个“层”（layer），模型都会存储矩阵 $W_Q, W_K, W_V in mathbb{R}^{d imes d}$，并按以下方式计算“注意力”：
$$Q = X cdot W_Q, quad K = X cdot W_K quad ext{和} quad V = X cdot W_V$$ $$ ext{Attention}(Q,K,V) = ext{softmax}(Q K^T / sqrt{d}) cdot V$$
或者，用 Python 代码表示：
def attention(X, W_q, W_k, W_v):
Q = X @ W_q
K = X @ W_k
V = X @ W_v
Q_KT = Q @ K.transpose(2, 1) # 注意：这里假设输入是批处理的
return softmax(Q_KT / sqrt(d_model)) @ V其中 $ ext{dot-product}$ 是两个矩阵的点积。
在实际操作中，多个 LLM 对话是并行处理的。因此，推理是批处理 (batched) 的——我们同时处理 $B$ 个聊天会话。这意味着我们的输入序列 $X in mathbb{R}_{B imes N imes d}$。
请在纸上推导一下数学过程来验证它是否匹配。在我们的 Python 代码中，只需要改变转置的参数：
原始 (单次序列)： $ ext{Q\_KT} = ext{Q} @ ext{K}. ext{transpose}(2, 1)$ 批处理 (Batch)： $ ext{Q\_KT} = ext{Q} @ ext{K}. ext{transpose}(0, 2, 1)$
然而，我们的注意力机制实现有一个问题：它从内存中读取了过多的数据。
让我们来看一次矩阵乘法 $K = X cdot W_k$。
（注：原文在此处被截断，但根据上下文，接下来的内容将讨论如何优化这个矩阵乘法，以适应服务高达 200k 个词元的聊天会话。）
原文出处：https://injuly.in/blog/napkin-inference-cost/index.html
