杨国强-以卓越性价比释放开放大模型潜能：TPU 上的推理优化全解 conv.docxVIP

以卓越性价比释放开放大模型潜能TPU上的推理优化全解

杨国强

：谷歌云/AIInfra架构师

目录

01TPU简介与版本演进

TPU的演进

02在TPU上部署开放模型的关键技术

推理架构图(框架:vllm/Max*/JetStream/Pathways)

InferenceServer

UserApplication

Query

HTTP/gRPCrequest

HTTP/gRPCresponse

PyTorchBasedLLM(llama,DeepSeek)

JAXBasedLLM(Gemma)

Diffusion

fvllm

g

Jetstream(MaxText)

sMaxDiffusion

Metrics

Hardware(TPU)

vllm简介

●目标:

○ 优化KV缓存内存利用率，实现高吞吐量模型服务●洞见:

○ 使用虚拟内存分页以减少内存浪费

vllmonTPU时间线ExecutorPerformance

Prefill Decode

Bottleneck

Optimization

Benefit新特性:

Prefilltokenredundancy

Multi-queryAttention(TPU-friendly)

Prefixcaching

Available

SlowKVcachewrites

Sparsecoreoffload+CacheLayoutRedesign

Fasterprefill/decode

Available

Noparallelisminprefill/decode

RaggedPagedAttention(TPU-friendly)

Chunkedprefill

Available

Multi-Lora

Available

Multimodalmodelsupport

Available

+PrometheusMetrics

Available

V1refactor

Available

Quantizedmodels

Available

APC(自动前缀缓存)

重复利用先前生成的计算结果，消除冗余处理。缓存在分层的内存中(HBM/主机内存)。

针对常见提示词前缀的重复请求减小TTFT提高吞吐量并最终达到降本增效的目的

典型的应用场景:多轮对话，针对技术文档/书籍或者视频的重复问答

缓存命中的业务场景下显著提升系统吞吐量

/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_advanced_features.ipynb

分块预填充(ChunkedPrefill)

连续批处理(ContinuousBatching)

03TPU硬件特性—和GPU有什么不同

v6epod支持万亿规模参数模型

Google’shigh-speedinterconnectmakesCloudTPUPodsAISupercomputer

All-reduce

byhardware

Aseasyasusing

asinglecomputer

VersatileandhighperformancemodelservingRangeofshapestomatchmodelsizeandcomplexity

16chips

16chipsVMShape Chips/Slice S

16chips

1 1

2x2 4

2x4 8

4x4 16

4x8 32

8x8 64

8x16 128

16x16 256

13BparamLLM

32.5BparamLLM

65BparamLLM

170BparamLLM

280BparamLLM

540BparamLLM

1TparamLLM

2TparamLLM

*GoogleInternalData.August2023.Batchsize=1.Multi-headattentionbaseddecoderonlylanguagemodels:prefixlength=2048,decodesteps=256,beamsize=32forsampling

OptimizedonvLLMfromDay1

Ironwood

7thGeneration