当大型语言模型生成单个 token 时实际发生了什么

📌 One-Sentence Summary

这篇文章一步步演示了一个解码器-only 的大型语言模型如何生成一个 token，揭示解码阶段受内存带宽限制，GPU 大约在 98%的时间等待权重传输。

📝 Summary

本文追踪了一个现代解码器-only Transformer（如 LLaMA 3 或 Mistral）中单个 token 的精确旅程。它首先对比计算受限的 prefill 阶段，其中所有 prompt token 并行处理，与受内存带宽限制的 decode 阶段形成对比，后者由于自回归特性每次只能生成一个 token。接着详细介绍了嵌入查找、RoPE 位置编码、包含 RMSNorm 的 32 层 transformer、Grouped-Query Attention 和 KV 缓存、SwiGLU 前馈网络、LM 头产生 128，256 个 logits，以及温度、top-K、top-P、softmax 和分类抽样采样管道。最后量化了 RTX 4090 上的硬件现实：生成一个 token 需要通过 1，008 GB/s 总线流式传输 16 GB 的权重，产生每秒约 63 个 token 的速度，且仅有~1.2%的计算利用率。文章结尾解释了量化、连续批处理和推测解码如何攻克这一内存墙。

💡 Main Points

Prefill 是计算受限，而解码是受内存带宽限制，这解释了快速 prompt 处理与缓慢逐 token 生成之间的不对称。

在 prefill 阶段，一个[T x d_model]矩阵乘以权重，使得每个加载的权重被所有 T 个 token 重复使用，从而饱和张量核心。在 decode 阶段，输入是一个[1 x d_model]向量，每个权重都必须从 VRAM 获取，只需两次 FLOPs，导致计算资源闲置。

单个解码步骤涉及嵌入查找、RoPE 旋转、含 RMSNorm 的 32 层 transformer、GQA 以及 KV 缓存、SwiGLU MLP，以及输出 128，256 个 logits 的 LM 头。

单个解码步骤包括嵌入查找、RoPE 旋转、含 RMSNorm 的 32 层 transformer、GQA 以及 KV 缓存、SwiGLU MLP，以及输出 128，256 个 logits 的 LM 头。

温度、top-K、top-P、softmax 和伪随机抽样的采样管道将原始 logits 转换为单个选定的 token。

温度、top-K、top-P、softmax 和伪随机抽样的采样管道将原始 logits 转换为单个选定的 token。

8B FP16 模型占用 16 GB。当速率为 1，008 GB/s 时，读取全部权重大约需要 15.87 毫秒，将生成速度限制在每秒约 63 个 token。实际计算速率约为 1.008 TFLOPS，对比 82.6 TFLOPS 的容量，利用率约为 1.22%。

8B FP16 模型占用 16 GB。当速率为 1，008 GB/s 时，读取全部权重大约需要 15.87 毫秒，将生成速度限制在每秒约 63 个 token。实际计算速率约为 1.008 TFLOPS，对比 82.6 TFLOPS 的容量，利用率约为 1.22%。

量化、连续批处理和推测解码都是针对内存带宽瓶颈的直接应对措施。

量化、连续批处理和推测解码都是针对内存带宽瓶颈的直接应对措施。

💬 Key Quotes

当你从大型语言模型中流式传输响应并观察单词一个接一个时，你的 GPU 正处于一种奇异的物理现实中：它大约花费 98%的时间等待内存传输，少于 2%的时间进行实际数学运算。

你的 GPU Tensor Cores 大约在 98.7%的时间完全空闲，等待权重通过内存总线传输。

它不以词语思考：它将整数 token ID 转换为 4096 维几何坐标。

它输出概率而非确定性：最终层产生 128，000 个原始 logits，这些 logits 经过过滤、温度缩放和采样。

它受限于内存速度而非计算能力：每个 token 都需要在 GPU 内存总线上扫描模型权重的所有字节。

📊 Article Meta

AI Screening： 86

Source： DEV Community: machinelearning

Author： Syed Anzar

Category： 人工智能

Language： 英文

Read Time： 9 min

Word Count： 2097

Tags：

AI 与智能应用 , LLM 推理优化 , 模型训练与推理 , KV Cache优化 , 大语言模型 (LLM)