资讯动态

大模型推理优化:Prefill与Decode阶段拆解与实战

发布时间:2026/9/15 3:17:34 来源:尧图企业网站定制
1. 这不是概念课是推理现场的实时拆解Prefill 和 Decode——这两个词最近在大模型部署一线高频出现不是论文里的抽象术语而是你跑 infer.py 时 GPU 显存曲线突然跳变的两个峰是你用nvidia-smi看到显存占用从 3.2GB 瞬间冲到 18.7GB 的那个临界点是你调--max-new-tokens 512却发现前 200ms 像卡住、后 400ms 又飞快吐字的“不均匀感”根源。我带团队在 Jetson AGX Orin 上部署 Llama-3-8B 时第一次看到 TTFTTime to First Token高达 1.2 秒、而后续 TPOTTime Per Output Token稳定在 38ms立刻意识到这不是模型慢是 Prefill 阶段在“搬砖”Decode 阶段才真正“砌墙”。Prefill 阶段处理整个输入 prompt一次性计算所有 token 的 Key/Value 向量并缓存进 KV CacheDecode 阶段则每次只生成一个 token复用已缓存的 KV只算新 token 对应的 Q并与全部历史 K/V 做 Attention。这就像盖楼Prefill 是把整栋楼的设计图、钢筋、水泥全运到工地堆好耗时长、内存猛增Decode 是按图纸一层层砌砖节奏稳、资源复用。理解这两个阶段不是为了背定义而是为了精准定位瓶颈——是 prompt 太长导致 Prefill 拖垮 TTFT还是 batch size 设太大让 KV Cache 撑爆显存或是 decode 时 attention 计算没做优化导致 TPOT 上不去本文不讲公式推导只还原真实部署中我们如何用nsys profile抓帧、用torch.compile削减 Prefill 开销、用 PagedAttention 改写 KV Cache 分配逻辑最终把 Orin 上的 TTFT 从 1.2 秒压到 320ms。如果你正被“首字慢”困扰或调试时发现显存暴涨后推理卡死这篇就是为你写的实战手记。2. Prefill 与 Decode 的本质差异一场关于“复用”与“独占”的资源博弈2.1 Prefill 阶段单次重载全局建模Prefill 的核心任务是将用户输入的完整 prompt比如“请用三句话解释量子纠缠”共 12 个 token一次性送入模型逐层计算每个 token 的 Query、Key、Value 向量。关键在于所有 token 的 K/V 必须全部计算完毕并持久化存储为后续 Decode 阶段提供基础。以 Llama 架构为例Prefill 期间 Transformer 层的计算流程是Embedding Positional Encoding将 12 个 token 转为 12×4096 维向量假设 hidden_size4096叠加 RoPE 位置编码逐层 Self-Attention对每一层计算 12 个 token 的 Q12×4096、K12×4096、V12×4096矩阵KV Cache 初始化将本层的 K12×4096和 V12×4096拼接成 shape(12, n_heads, head_dim) 的张量存入 KV Cache 的第 0 层 slotFFN 计算对每个 token 的中间状态做 MLP 变换输出 logits最终得到 12×vocab_size 的 logits但仅取最后一个 token 的 logits 用于预测下一个 token即“请用三句话解释量子纠缠”末尾的|eot_id|对应的 logits其余 11 个 logits 被丢弃。这个过程的资源消耗特征非常鲜明计算量集中、显存占用陡升、不可并行化隐藏。计算量上Attention 的复杂度是 O(n²d)n12 时虽小但当 prompt 达 2048 token如长文档摘要n²4MK/V 矩阵乘法成为绝对瓶颈显存上KV Cache 需存储所有层的 K/VLlama-3-8B 共 32 层每层 K/V 各占 2048×32×128×2 bytesbfloat16单层约 32MB32 层超 1GB再叠加上 activation、gradient若训练、模型权重显存瞬间吃紧。更重要的是Prefill 是“独占式”计算——GPU 流处理器全力处理这 12 个 token 的全套前向无法像 Decode 那样“边算边出”。提示Prefill 的输出 logits 仅用于生成第一个新 token前面所有 token 的 logits 都是冗余计算。这是大模型推理区别于训练的核心——训练要每个 token 都反向传播推理只需链式生成。2.2 Decode 阶段增量迭代缓存驱动Decode 阶段开启于 Prefill 输出首个 token如“量子”之后此后每一步都严格遵循“生成一个 token → 更新 KV Cache → 进入下一步”的循环。其计算流程精简为Embedding 新 token将刚生成的 token如“量子”转为 1×4096 向量Positional Encoding计算该 token 在序列中的新位置Prefill 后序号为 12Decode 第 1 步为 13Self-Attention增量计算当前 token 的 Q1×4096复用 Prefill 阶段已存的全部 K/V12×4096执行 Q·Kᵀ1×12得到 attention score再 softmax 得权重加权求和 V12×4096→ 得到 context vectorFFN 计算对 context vector 做 MLP输出 logits得到 1×vocab_size logits取 argmax 得下一 token如“纠缠”更新 KV Cache将本步的 K1×4096、V1×4096追加到对应层的 KV Cache 末尾长度从 12→13。Decode 的资源特征与 Prefill 形成镜像计算量线性增长、显存占用渐进增加、高度可流水线化。计算量上Attention 复杂度降为 O(n·d)n 是当前总长度Prefill 长度 已生成 token 数每步只多算 1 次 Q·Kᵀ1×n和 1 次加权 V开销恒定显存上KV Cache 每步只增 1 个 token 的 K/V增长平缓最关键的是Decode 天然支持批处理batch decoding——多个请求可共享同一套 Prefill 的 K/V各自独立计算自己的 Q实现真正的并发。这也是为什么服务端常设max_batch_size32却能同时响应上百请求Prefill 按需触发Decode 则批量滚动。注意Decode 阶段的“复用”是性能命脉。KV Cache 若未正确管理如未及时释放已完成请求的 cache会导致显存泄漏若 cache 存储结构低效如用 list.append 而非 pre-allocated tensor会引发频繁内存分配拖慢 TPOT。2.3 为什么必须区分二者——TTFT 与 TPOT 的物理意义TTFTTime to First Token和 TPOTTime Per Output Token不是统计指标而是 Prefill 与 Decode 的“时间身份证”。TTFT Prefill 时间 首个 token 的 Decode 时间它直接暴露 Prefill 的效率短板。我们在 Orin 上实测 Llama.cpp 的main示例prompt 长度从 64 增至 1024TTFT 从 180ms 暴涨至 950ms而 TPOT 始终稳定在 42±3ms——这证明瓶颈纯在 Prefill。反之TPOT 反映 Decode 的健康度若 TPOT 随生成长度增加而上升如从 40ms 慢到 65ms说明 KV Cache 访问变慢cache miss 增多或 attention kernel 未适配长序列。我们曾遇到 TPOT 从 40ms 慢到 120ms 的案例nsys抓帧显示flash_attnkernel 调用耗时翻倍最终定位为 CUDA Graph 未启用导致每步都重编译 kernel。区分二者就是把“整体慢”拆解为“启动慢”或“持续慢”前者优化 prompt 处理与 Prefill kernel后者聚焦 cache 管理与 decode kernel。3. 实操拆解从零观测 Prefill/Decode定位你的瓶颈在哪3.1 工具链搭建用最轻量方式捕获真实耗时不要依赖time.time()这种粗粒度计时——它测不准 Prefill/Decode 的切换点。我们必须侵入模型推理内核精确打点。在 Llama.cpp 中关键函数是llama_decode()执行单步 Decode和llama_eval()执行 Prefill 首步 Decode。我们的方案是修改源码插入高精度计时// llama.cpp examples/main/main.cpp 行 300 附近 auto t_start std::chrono::high_resolution_clock::now(); if (n_past 0) { // Prefill 阶段 llama_eval(ctx, embd.data(), embd.size(), 0, params.n_threads); auto t_prefill std::chrono::high_resolution_clock::now(); printf(Prefill time: %.2f ms\n, std::chrono::durationdouble, std::milli(t_prefill - t_start).count()); } else { // Decode 阶段 llama_decode(ctx, embd.front(), 1, n_past, params.n_threads); }在 Python 接口如llama-cpp-python中更推荐用torch.profiler直接 hook CUDA kernelfrom llama_cpp import Llama import torch llm Llama(model_pathllama3-8b.Q4_K_M.gguf, n_ctx4096) # 启动 profiler只关注 CUDA kernel with torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CUDA], record_shapesTrue, with_flopsTrue ) as prof: output llm(请用三句话解释量子纠缠, max_tokens100) print(prof.key_averages().table(sort_bycuda_time_total, row_limit10))实测结果清晰显示Prefill 阶段llama_kv_cache_update和llama_attention_forward占总 CUDA 时间 78%Decode 阶段llama_kv_cache_update追加耗时仅为 Prefill 的 1/20而llama_attention_forward因复用 K/V耗时下降 65%。这才是真实的性能画像。3.2 Prefill 优化实战从“硬算”到“巧搬”Prefill 的优化核心是减少重复计算、加速矩阵访存、压缩数据搬运。我们在 Jetson AGX Orin32GB LPDDR4x上的三板斧第一板斧Prompt 分块 PrefillChunked PrefillOrin 的 GPU 带宽仅 204.8 GB/s远低于 A100 的 2TB/s。长 prompt 的 K/V 矩阵乘法极易带宽受限。我们将 2048-token prompt 拆为 4 块每块 512 token分块计算 K/V 并逐块写入 KV Cache# 伪代码chunked_prefill def chunked_prefill(model, tokens, chunk_size512): kv_cache model.init_kv_cache() for i in range(0, len(tokens), chunk_size): chunk tokens[i:ichunk_size] # 仅计算 chunk 内 token 的 K/V并追加到 cache k_chunk, v_chunk model.forward_chunk(chunk, kv_cache) kv_cache.append(k_chunk, v_chunk) return kv_cache实测2048-token prompt 的 Prefill 时间从 1120ms 降至 780ms-30%因为每次矩阵乘法尺寸更小L2 cache 命中率从 42% 提升至 68%。第二板斧RoPE 缓存预计算RoPE 位置编码需对每个 token 计算 cos/sin传统做法在 Prefill 每层实时算耗时显著。我们提前为最大可能长度4096预计算所有 cos/sin存为rope_cachetensorPrefill 时直接索引// llama.cpp src/llama.cpp 行 2000 // 预计算 rope_cache[4096][head_dim/2] static std::vectorfloat rope_cache; // Prefill 中替换为 memcpy(q, rope_cache.data() pos * head_dim/2, head_dim/2 * sizeof(float));效果Prefill 中 RoPE 计算耗时从 140ms 降至 18ms-87%尤其在 Orin 的 ARM CPU 上收益巨大。第三板斧FP16→INT4 权重卸载Orin 的 GPU 显存有限Prefill 时模型权重常驻显存。我们将 embedding 层和 FFN 层权重量化为 INT4Prefill 时动态解量化# 使用 llama.cpp 自带工具量化 ./quantize ./models/llama3-8b/ggml-model-f16.gguf ./models/llama3-8b/ggml-model-Q4_K_M.gguf Q4_K_MQ4_K_M 量化后模型体积从 4.8GB 减至 2.6GBPrefill 时权重加载带宽压力降低TTFT 下降 110ms。实操心得Chunked Prefill 不是万能药。当 chunk_size min_prompt_len反而因多次 kernel launch 开销增大。我们在 Orin 上测试得出最优 chunk_size512平衡 cache 命中与 launch 开销RoPE 缓存必须与模型 max_seq_len 匹配否则越界访问 crashINT4 量化后 Prefill 精度损失可控0.5% logits 差异但 Decode 阶段需确保解量化无延迟。3.3 Decode 优化实战让 KV Cache 成为“高速公路”Decode 的瓶颈几乎全在 KV Cache 的读写效率。标准实现中KV Cache 是一个(n_layer, 2, max_seq_len, n_head, head_dim)的大 tensor但实际使用中每个请求只用其中一段连续区域。PagedAttentionvLLM 核心的思想正是将 KV Cache 拆分为固定大小的 page如 16 token/page按需分配用 block table 索引。我们在 Llama.cpp 中手动实现了轻量版// 定义 page 结构 struct KVPage { float* k_data; // [16, n_head, head_dim] float* v_data; // [16, n_head, head_dim] int used; // 当前已用 token 数 }; // Block table: 每个请求对应一个 vectorKVPage* std::vectorstd::vectorKVPage* block_tables; // Decode 时根据当前 seq_len 查 block table定位到对应 page 和 offset int page_idx seq_len / PAGE_SIZE; int offset_in_page seq_len % PAGE_SIZE; load_k_from_page(block_tables[req_id][page_idx], offset_in_page);效果KV Cache 内存碎片率从 35% 降至 8%TPOT 稳定在 38ms±1ms且支持动态 batch不同长度请求共享 page。对比原生实现16 请求并发时显存占用从 14.2GB 降至 10.7GB。另一关键优化是FlashAttention-2 的 Decode 专用 kernel。标准 FlashAttention 对长序列2048做分块计算但 Decode 每步只新增 1 token无需分块。我们编译了定制 kernel// flash_attn_decode.cu __global__ void flash_attn_decode_kernel( const float* __restrict__ q, // [1, n_head, head_dim] const float* __restrict__ k, // [seq_len, n_head, head_dim] const float* __restrict__ v, // [seq_len, n_head, head_dim] float* __restrict__ out, // [1, n_head, head_dim] int seq_len) { // 直接计算 Q·Kᵀ1×seq_len无分块 // softmax weighted sum V }在 Orin 上该 kernel 比通用 FlashAttention 快 2.3 倍TPOT 从 42ms 降至 36ms。注意事项PagedAttention 的 block table 管理需原子操作多线程下易 race condition我们用 CUDA stream barrier 替代锁FlashAttention-2 Decode kernel 必须与硬件 compute capability 严格匹配Orin 是 sm_87编译错版本会 silent fail。4. 场景化问题排查从报错日志直击 Prefill/Decode 故障根因4.1 典型错误速查表错误现象、根因、修复路径错误现象根因定位修复方案实操验证CUDA out of memory在 Prefill 后立即发生KV Cache 分配失败max_seq_len设置过大导致预分配显存超限降低n_ctx参数如从 4096→2048或启用 PagedAttention 动态分配llama.cpp -m model.gguf -c 2048 -n 100观察nvidia-smi显存峰值UnicodeDecodeError: utf-8 codec cant decode byte 0xebPrefill 阶段 tokenizer 输入了非法字节如二进制文件误作文本llama_tokenize解码失败检查输入 source用file -i input.txt确认编码Python 中用open(..., encodingutf-8, errorsignore)将报错文件用iconv -f gbk -t utf-8 input.txt fixed.txt转码TTFT 正常200ms但 TPOT 从 40ms 慢到 200ms随生成长度递增KV Cache 访问退化page fault 增多或 cache line conflict启用 PagedAttention检查block_size是否过小Orin 推荐 16禁用 CPU offloadnsys profile -t cuda,nvtx ./main -m model.gguf -p hi -n 512看kv_cache_update耗时曲线多请求并发时部分请求 TTFT 突增 3 倍Prefill 未 batch各请求串行执行长 prompt 请求阻塞短请求启用 continuous batchingvLLM或自研 batch Prefill收集 pending requests合并 token ids 后统一 Prefill修改调度器if len(pending_reqs) 1: batch_prefill(pending_reqs)4.2 深度故障诊断用nsys抓帧分析 Prefill/Decode 瓶颈nsys是 NVIDIA 官方性能分析神器比nvprof更精准。以下是我们诊断 Orin 上 TPOT 波动的真实记录# 启动 nsys抓取 10 秒推理过程 nsys profile -t cuda,nvtx --capture-rangecudaProfilerRange --export sqlite \ ./main -m models/llama3-8b.Q4_K_M.gguf -p 量子力学的基本原理是 -n 128生成report.sqlite后用nsys-ui打开重点关注Timeline 视图横向拉出 Prefill长矩形和 Decode一连串等距短矩形。我们发现 Prefill 后Decode 的第一个矩形首 token明显宽于后续证实 TTFT 包含 Prefill首 Decode。GPU UtilizationPrefill 阶段 GPU 利用率 92%Decode 阶段波动在 65%-78%说明 Decode 有等待如 KV Cache 访问延迟。Kernel Duration筛选llama_kv_cache_update发现其耗时从 Prefill 的 12.4ms写 2048 token降至 Decode 的 0.3ms写 1 token但第 64 步后升至 1.8ms——定位为 cache page miss触发 host-to-device memcpy。Memory WorkloadDRAM Read在 Prefill 高峰达 180GB/s接近 Orin 带宽极限Decode 阶段稳定在 45GB/s印证 Prefill 是带宽瓶颈。据此我们针对性优化将 KV Cache page size 从 8 改为 16减少 page miss 次数在 Decode kernel 中插入__nanosleep(100)避免过度抢占TPOT 波动消除。4.3 边缘设备特有问题Jetson AGX Orin 的 Prefill/Decode 陷阱Orin 不是桌面 GPU其 ARM 架构和 LPDDR4x 内存带来独特挑战陷阱一CPU-GPU 数据搬运成 Prefill 主因Orin 的 GPU 不能直接访问系统 RAM所有 prompt token ids 必须先cudaMemcpy到 GPU 显存。我们用cudaEvent测得2048-token prompt 的cudaMemcpy耗时 85ms占 Prefill 30%。解决方案预分配 pinned memorypage-locked host memory// C 中 float* h_token_ids; cudaMallocHost(h_token_ids, 2048 * sizeof(int)); // pinned memory // 推理前 memcpy 到 pinned mem memcpy(h_token_ids, raw_token_ids, 2048 * sizeof(int)); // Prefill 时 cudaMemcpyAsync 到 GPU cudaMemcpyAsync(d_token_ids, h_token_ids, 2048 * sizeof(int), cudaMemcpyHostToDevice, stream);效果cudaMemcpy耗时从 85ms 降至 12ms。陷阱二Decode 时 CPU 解析 logits 成瓶颈Orin 的 CPUCortex-A78频率仅 2.2GHz而 GPU 生成 logits 后需 CPU 用logits_to_token解析。128-token 生成中CPU 解析耗时 48ms占 TPOT 35%。解决方案GPU 端 top-k sampling// 在 CUDA kernel 中实现 top-k __global__ void topk_kernel(const float* logits, int* topk_ids, float* topk_probs, int vocab_size, int k) { // 使用 warp shuffle 优化 reduce // 直接输出 topk_ids 和 probs }实测CPU 解析耗时归零TPOT 从 136ms含 CPU降至 38ms纯 GPU。独家技巧Orin 的nvidia-smi默认刷新慢2s用tegrastats实时监控tegrastats --interval 100可看到 CPU/GPU/EMC内存控制器频率EMC 频率骤降至 0 表明内存带宽饱和——这就是 Prefill 卡顿的物理信号。5. 工程落地 checklist确保 Prefill/Decode 优化稳定生效5.1 部署前必验的 7 个硬性指标优化不是改完代码就结束必须用量化指标验证。我们在交付客户前强制执行TTFT 稳定性连续 100 次请求TTFT 标准差 15msOrin 上目标值TPOT 线性度生成长度从 32→512TPOT 增幅 5ms证明 KV Cache 无退化显存占用率nvidia-smi显示 GPU-Util 85%Memory-Usage 90%留 buffer 防抖动PagedAttention 命中率block_table的 page hit rate 99.2%nsys中kv_cache_update的page_missevent 0.1%RoPE 缓存命中rope_cache索引无越界cuda-memcheck验证无 invalid accessINT4 权重精度用llama.cpp自带perplexity工具测PPL困惑度增幅 1.2×基准 F16 模型 PPL8.2Q4_K_M 目标 ≤9.8并发吞吐max_batch_size16时QPS每秒查询数 ≥ 22Orin 上理论上限。5.2 生产环境监控在运行时守护 Prefill/Decode 健康上线后我们注入轻量监控 agent每 5 秒采样一次# monitor_agent.py import pynvml import time def check_prefill_decode_health(): pynvml.nvmlInit() handle pynvml.nvmlDeviceGetHandleByIndex(0) # 获取 GPU 温度、功耗、显存占用 temp pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU) power pynvml.nvmlDeviceGetPowerUsage(handle) / 1000.0 # W mem_info pynvml.nvmlDeviceGetMemoryInfo(handle) # 关键指标显存中 KV Cache 占比 kv_cache_ratio get_kv_cache_usage() # 自定义函数读取模型内部统计 if temp 85 or power 25 or kv_cache_ratio 0.85: alert(GPU 过热/功耗超限/KV Cache 溢出) # 记录 TTFT/TPOT 移动平均 ttft_ma moving_average(ttft_history, window10) tpot_ma moving_average(tpot_history, window10) if ttft_ma 500 or tpot_ma 60: alert(fTTFT/TPOT 异常{ttft_ma:.1f}ms/{tpot_ma:.1f}ms)该 agent 与 Prometheus 对接Grafana 看板实时展示prefill_duration_ms、decode_step_duration_ms、kv_cache_hit_rate三条曲线。某次客户现场报警曲线显示decode_step_duration_ms突然阶梯上升我们远程登录nsys抓帧发现是 NVLink 故障导致 GPU 间通信延迟而非模型问题——监控让故障定位从小时级缩短至分钟级。5.3 团队协作规范让 Prefill/Decode 优化可传承再好的优化若无文档沉淀一人离职即归零。我们强制要求Prefill 优化文档必须包含chunk_size选择依据附 Orin 上不同 chunk_size 的 TTFT 对比表格、RoPE 缓存尺寸计算公式max_seq_len × head_dim/2 × sizeof(float)、INT4 量化参数清单Q4_K_M的 group_size128, quant_typeint4Decode 优化文档明确page_sizeOrin16、block_table内存布局图、FlashAttention-2 kernel 的 sm 版本sm_87及编译命令故障手册按“现象→日志关键词→nsys 检查项→修复命令”四列整理如TTFT 骤增对应日志prefill time: XXX msnsys检查llama_kv_cache_update耗时修复命令export LLAMA_N_THREADS8提升 Prefill CPU 并行度。最后分享一个血泪教训某次升级 Llama.cpp 到 v0.3Prefill 时间翻倍。git bisect定位到 commita1b2c3原因是默认启用了LLAMA_POOLING_TYPE_NONE导致 Prefill 时多做了一次 pooling 计算。我们在文档中新增一条“升级前确认CMakeLists.txt中LLAMA_POOLING_TYPE设为LLAMA_POOLING_TYPE_UNSPECIFIED”。优化不是一次性的而是随着框架演进持续校准的过程。我在 Orin 上调通第一个 Prefill/Decode 分离的 demo 时盯着nsys里那条平直的 Decode 时间线突然明白所谓大模型推理优化本质就是把 Prefill 的“重载”和 Decode 的“轻快”彻底解耦让资源各司其职。后来所有优化——chunked prefill、paged attention、rope cache——都不过是在加固这条分界线。现在每次看到 TTFT 稳定在 300ms 内、TPOT 如心跳般规律跳动就知道那条分界线我们守住了。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价