TPU优化技巧¶
本文档汇集了优化您在TPU上运行vLLM工作负载的实用技巧。
快速开始¶
正在寻找安装和配置指南？请点击此处查看。
TPU工作负载规模调整¶
在选择单个服务实例的理想芯片数量时，必须同时考虑模型大小和平均请求上下文长度。确保KV缓存有足够的HBM内存对于处理充足的并发请求量至关重要。
以下colab 计算器将告诉你：
- 每个token和每个请求的KV缓存大小需求
- 模型权重占用的TPU/GPU内存
- 为KV缓存分配的TPU/GPU内存
- 可设置的大致最大请求数量 (--max-num-seqs)
这种方法可作为一般经验法则。
延迟与吞吐量的权衡¶
与根据工作负载调整芯片数量类似，可以考虑调整--max-num-seqs参数来微调延迟与吞吐量之间的平衡。减少--max-num-seqs值和/或增加芯片数量有助于降低延迟。
--max-num-seqs 定义了并发解码槽的数量，实际上限制了服务器可以同时处理的请求令牌数。增加该值允许服务器预分配更多HBM来处理更高数量的并发请求，从而最大化整体吞吐量。然而，这通常会增加每个请求的端到端(e2e)延迟。
因此，仔细调整--max-num-seqs参数对于在特定工作负载下实现延迟和吞吐量之间的理想平衡至关重要。
类似地，可以调低--max-num-batch-tokens来改善延迟，或调高来提升吞吐量。
编译与缓存¶
对于有GPU背景的用户来说，TPU最显著的区别之一在于初始编译步骤。TPU是专用加速器(ASIC)，通过XLA编译器执行预编译的静态计算图来实现最佳性能。与能更灵活处理动态输入形状的GPU不同，TPU需要为每个处理的张量形状(如批大小和序列长度)编译特定的计算图。
为了管理这一点，vLLM在首次启动服务器时会执行一次"预热"过程。在此阶段，它会为各种常见输入形状预编译模型，并将这些编译好的计算图保存到磁盘或远程存储的缓存中（默认位于~/.cache/vllm/xla_cache）。这个过程的时间差异很大，从几分钟到一小时不等，具体取决于模型大小和使用的上下文长度。
虽然首次编译可能需要一些时间，但对于后续所有服务器启动，vLLM可以直接从缓存加载这些计算图，从而消除后续运行的编译时间。
使用VLLM_XLA_CACHE_PATH环境变量将数据写入可共享存储，供未来部署的节点使用（例如在使用自动扩展时）。
减少编译时间¶
初始编译时间差异很大，并受到本优化文档中讨论的许多参数影响。影响编译时长的因素包括模型大小和--max-num-batch-tokens等。您可以调整的其他参数包括VLLM_TPU_MOST_MODEL_LEN等。
基于您的数据进行优化¶
最大模型长度 vs. 大多数模型长度¶
如果您的请求大多短于模型最大长度，但仍需处理偶尔出现的较长请求，设置过高的最大模型长度可能会影响性能。这种情况下，您可以通过指定VLLM_TPU_MOST_MODEL_LEN环境变量来尝试引入most model len。
例如，1%的请求长度为32k，99%的请求长度为2k。您可以将32k传入--max-model-len 32768并使用VLLM_TPU_MOST_MODEL_LEN=2048。
请求被细分为max-model-len和most-model-len两类，对于后者，我们可以获得更好的性能，因为服务器可以同时处理更多请求。
填充¶
对于有延迟要求的在线服务，考虑通过设置环境变量VLLM_TPU_BUCKET_PADDING_GAP切换到桶填充方式。由于TPU的布局特性，建议尝试以128为增量进行设置：128、256等。
服务器在将请求发送到模型之前会将其填充为固定长度，以避免重新编译。要了解更多关于TPU填充的信息，请参阅此处。目前有两种填充请求的方式：
1) 默认的指数填充（填充到最接近的2的幂次） 2) 桶填充（填充到最接近的线性递增桶）。
使用桶填充时，桶的起始大小为16，结束于max_model_len，并以VLLM_TPU_BUCKET_PADDING_GAP为增量递增。
例如，max_model_len=512，padding_gap=64时，分桶大小将为[16, 32, 64, 128, 192, 256, 320, 384, 448, 512]。
我们填充的token越少，TPU执行的不必要计算就越少，性能就能越好。例如，如果num_tokens=300，采用指数填充会填充到512，而使用上述bucket_padding则只填充到320。
不过，选择填充间隙时需要谨慎。如果间隙过小，意味着桶的数量会增多，从而导致预热（预编译）时间延长，并需要更多内存来存储编译后的计算图。过多的编译图可能导致HBM内存溢出。反之，若间隙过大，与默认的指数填充相比则无法获得性能提升。
量化¶
如果可能，请使用与芯片硬件加速相匹配的精度：
- v5e在MXU中支持int4/int8硬件加速
- v6e在MXU中支持int4/int8硬件加速
vLLM在TPU上支持的量化格式与功能 [2025年7月]：
- INT8 W8A8
- INT8 W8A16
- FP8键值缓存
- [开发中] FP8 W8A8
- [开发中] AWQ
- [开发中] FP4 W4A8
并行化¶
在单机部署时，不要将TP设置为小于芯片数量的值。
虽然通常使用GPU来实现这一点，但不要尝试在单台主机的8个芯片上分散运行2或8个不同的工作负载。如果您需要1或4个芯片，只需创建具有1或4个芯片的实例（这些是部分主机机型）。
调优您的工作负载¶
尽管我们努力提供出色的默认配置，但我们强烈建议您查看vLLM自动调优器，以便根据您的使用场景优化工作负载。
我们将涵盖的未来主题¶
性能分析¶
自动调优器在最后一步会提供一份优化配置的分析报告。然而，对于新用户来说，解读这份报告可能具有挑战性。我们计划在未来扩展这一部分，提供更详细的指导。在此期间，您可以学习如何使用vLLM内置的分析工具此处收集TPU分析报告。这份报告可以为您的工作负载性能提供有价值的洞察。
SPMD¶
更多详情即将公布。
希望我们涵盖本文档未列出的内容？请提交一个issue并引用此文档。我们很乐意听取您的问题或建议。
