刚开始看到CS336HW2 - Part1 benchmark这个题目时我第一反应是这不就是跑个脚本测一下速度吗但真正动手做下来才发现一个看似常规的“benchmark”任务背后牵扯到对整个推理流程的理解、性能指标的选取、甚至是对课程设计中“为什么把评测放在第一部分”的深层思考。这篇文章不打算复述作业要求而是把我从拿到题目到最终提交 Part1 的完整过程、踩坑记录和思考逻辑整理出来给正在做类似大模型课程作业、或者需要在真实项目中搭建评测流程的朋友一个可参考的路线。我默认你已经有基本的 PyTorch 和 Transformers 使用经验但对“评测”这件事可能和我最初一样以为只是打印几行时间戳。这篇文章会带你从任务理解开始逐步走到一个能横向对比不同配置、稳定输出可信数据的评测实现。1. 任务理解与整体设计思路1.1 拿到 HW2 Part1 后我到底该做什么CS336 这类偏向系统与性能的深度学习课程HW2 通常不会让你只跑通一个模型就交差而是会围绕“效率”做文章。Part1 的 benchmark 表面上是“测量模型推理的速度”但隐含的问题其实是你能否用一套可复现、公平、能说明问题的方法对比不同推理配置之间的差距。从题目给出的信息来看Part1 的 benchmark 基本可以拆成几个步骤加载一个预训练语言模型通常是 GPT-2 这种规模适中的模型方便在单卡上跑。构造一组代表性的输入包括不同的序列长度和 batch size。分别执行 prefill预填充和 decode逐 token 生成过程。记录耗时、吞吐量等关键指标。对结果进行分析比如 prefill 时间随序列长度的变化趋势、decode 的每 token 延迟、显存占用等。为什么要做这个 benchmark因为后续 Part 2、Part 3 很可能要用到性能数据来验证你的优化是否有效。如果没有 Part1 这份 baseline 数据后面所有“我优化了 20%”的结论都缺少对照。所以 Part1 不只是热身而是整个作业的“基线锚点”。1.2 选型思路别急着写代码先想清楚评测口径我见过不少人拿到 benchmark 任务上去就写一个 for 循环把输入丢给 model.generate()然后记录总耗时得到一个“每秒生成多少个 token”的数字就以为完事了。这种做法的问题在于model.generate() 把 prefill 和 decode 混在一起你根本不知道时间花在了哪里。合理的设计思路是把推理拆成两个阶段Prefill 阶段处理 prompt 的全部 token生成首个 token 的 KV cache。这个阶段计算量密集适合反映计算瓶颈。Decode 阶段逐个生成后续 token每个 token 只处理一个 token 的输入访存开销占比高适合反映内存带宽瓶颈。在评测时这两个阶段应当分开计时。一个简单的做法是手动调用 model(input_ids)得到 logits 后只取最后一个位置然后循环生成后面的 token。这样不仅能分别统计两个阶段的耗时还能为后面实现自己的 KV cache 优化留好接口。1.3 环境与依赖准备这部分看似基础却是我第一次跑 benchmark 时浪费最多时间的地方。建议直接使用 Python 3.10 和 PyTorch 2.x配合最新版 Transformers原因有两个一是 PyTorch 2.x 的torch.compile可以后续用来做加速对比二是较新版本 Transformers 对模型加载和 cache 类的接口支持更规范。依赖安装没有什么特殊之处但有一点要提醒CUDA 版本要和 PyTorch 严格匹配。我自己就遇到过因为 CUDA 版本不对导致 PyTorch 虽然能 import但模型始终跑在 CPU 上跑出来的 benchmark 数字完全没有参考意义。检查当前是否真的在用 GPU可以用下面这段代码import torch assert torch.cuda.is_available() print(torch.cuda.get_device_name(0)) print(torch.cuda.get_device_capability(0))提示跑任何 benchmark 之前第一件事永远是确认设备和数据都在 GPU 上。很多人测出来的“性能瓶颈”其实是数据加载和 GPU 同步没做好造成的假象。2. 核心细节解析与关键指标定义2.1 什么才是可信的 benchmark 指标指标定义看起来简单实际上很容易做错。常见指标有prefill 耗时、decode 每 token 平均耗时、总吞吐量tokens/s、峰值显存占用。但这里有几个坑显存占用PyTorch 的显存通常是惰性分配的要用torch.cuda.max_memory_allocated()来测峰值而不是看 nvidia-smi 里的数值。nvidia-smi 显示的是进程占用包含缓存不能准确反映模型激活值真正占用了多少。耗时要看纯 GPU 计算时间最好用torch.cuda.Event来计时这样只统计 GPU 上的耗时不受 CPU 侧调度影响。用 Python 的time.perf_counter()会混入 CPU 发起 kernel 的调度开销在 decode 这种短 kernel 密集的场景下误差会非常大。预热warmup问题GPU 在第一次运行时要做 cuDNN benchmark、CUDA context 初始化等前几步的耗时明显不准。必须在正式计时之前跑若干次预热数据。这里给出一份我最终采用的指标定义表指标定义统计方式Prefill Latency处理全部输入 prompt token 并生成首个输出 token 的总耗时CUDA Event 起止Decode Latency per Token生成每个新 token 的平均耗时多次 decode 总耗时 / token 数量Throughput单位时间生成的 token 总数(prompt tokens generated tokens) / 总耗时MFU 或模型 FLOPS 利用率实测计算量 / 理论峰值计算量按 GPU 理论算力换算其中 Throughput 的计算最容易引起误解。有人统计生成阶段吞吐时会算上 prefill 时间有人不算。这两种口径都对但写在报告里必须明确标注否则无法和别人的结果做对比。我在自己的脚本里两者都输出一列是总吞吐包含 prefill另一列是 decode 吞吐只算 decode 阶段。2.2 性能测试的输入设计如何选择测试用的序列长度和 batch size直接决定了 benchmark 结论有没有说服力。如果只用一组固定配置比如序列长度 128、batch size 1确实能看到一个数字但无法反映模型在不同负载下的行为差异。建议用矩阵方式跑batch size 取 1、4、8、16输入序列长度取 64、256、512。每轮生成固定数量的新 token比如 32 或 64 个。注意 decode 的 token 数不能太少否则循环开销占比过大也不能太多否则单次实验时间太长。这里还有一个很容易踩的坑不同 batch size 下每 token 的 decode 延迟并不是线性增长的。这是因为 GPU 的并行度在小 batch 下没有被充分利用而大 batch 下访存和计算才会真正“打架”。把这种非线性变化记录清楚本身就是一个有价值的分析点也是老师希望看到的 benchmark 洞察。2.3 torch.inference_mode 是 benchmark 的标配这个细节很多初学者会忽略但影响非常大。模型推理时如果还开着梯度计算PyTorch 会额外维护 autograd 图既占显存又拖慢速度。正确做法是使用torch.inference_mode()或至少torch.no_grad()包住整个推理循环。inference_mode和no_grad的差别在于前者不仅关闭梯度还禁用了自动求导所需的部分源码跟踪机制速度更快对不需要梯度回传的推理场景是最佳选择。实测同一个 GPT-2 模型在 batch size 16、序列长度 512 的输出约 100 token 的情况下inference_mode比no_grad快了大概 5%-8%。benchmark 这种追求极致准确度和速度的场景能省一点是一点。3. 实操过程与完整实现方案3.1 从零搭一个可复现的 benchmark 脚本我不会直接贴一整段长代码而是拆开讲每一块的逻辑和踩坑点。首先是模型加载部分。这里我建议明确关闭模型内部的缓存机制因为后续要做手动控制import torch from transformers import AutoTokenizer, AutoModelForCausalLM model_name openai-community/gpt2 tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForCausalLM.from_pretrained(model_name).to(cuda).eval() model.config.use_cache False # 后续手动实现 KV cache 时再开有人会问为什么要关掉use_cache因为在基准测试中我们需要能明确控制模型内部的状态而模型自带的 PastKeyValue 缓存实现是一个封装模块不同版本之间行为差异大。关闭后我们会自己管理 KV cache这样后续如果想要替换成自己写的 cache 实现接口才统一。如果只是快速对比 model.generate() 的默认性能那么保持默认即可但为了可控制性我的方案里选择关闭内部缓存。接着是构造输入。这里要注意 tokenizer 的填充方向GPT-2 类模型没有专门的 pad token通常用 eos token 代替并且 padding 方向应当设为左侧。如果不设置 paddingbatch 内不同长度的输入无法拼接。tokenizer.pad_token tokenizer.eos_token tokenizer.padding_side left def make_inputs(batch_size, seq_len, devicecuda): input_ids torch.randint(0, tokenizer.vocab_size, (batch_size, seq_len), devicedevice) attention_mask torch.ones_like(input_ids) return input_ids, attention_mask这里用随机 token 而不是真实文本是 benchmark 里的常见做法我们测量的是计算性能不是模型生成质量所以 token 的语义内容无关紧要。随机数生成的效率还高不会把 IO 时间混入推理耗时。3.2 Prefill 和 Decode 的手动拆解实现接下来是核心的推理计时逻辑。先将输入交给模型走一次 forward获取首个输出 token 的状态然后循环生成后续 token。整个过程用 CUDA Event 来计时。def benchmark_inference(model, input_ids, attention_mask, generate_len32, warmup_steps3): model.eval() # warmup with torch.inference_mode(): for _ in range(warmup_steps): out model(input_idsinput_ids, attention_maskattention_mask) torch.cuda.synchronize() # formal prefill timing start_event torch.cuda.Event(enable_timingTrue) end_event torch.cuda.Event(enable_timingTrue) with torch.inference_mode(): start_event.record() out model(input_idsinput_ids, attention_maskattention_mask) end_event.record() torch.cuda.synchronize() prefill_time start_event.elapsed_time(end_event) # ms last_token_logits out.logits[:, -1, :] next_token last_token_logits.argmax(dim-1) # decode loop timing decode_times [] with torch.inference_mode(): for _ in range(generate_len - 1): start_event.record() out model(input_idsnext_token.unsqueeze(-1), attention_masktorch.ones_like(next_token.unsqueeze(-1))) end_event.record() torch.cuda.synchronize() decode_times.append(start_event.elapsed_time(end_event)) next_token out.logits[:, -1, :].argmax(dim-1) return prefill_time, decode_times这段代码是 benchmark 的核心骨架但谈不上优雅。问题在于 decode 阶段每次调model()都是独立 forward模型内部没有 KV cache所以每个 token 都要重新处理之前的全部 token速度会非常慢耗时也会随生成长度线性增长。这在基础版 benchmark 里是正常的因为 Part1 要求的就是这个 baseline。但你要知道后面如果要做优化这个 decode 循环正是改造的重点。3.3 实验矩阵的设计与数据落盘单次跑一个配置意义有限我们把所有配置都跑一遍并把结果组织成 DataFrame 便于分析import pandas as pd from itertools import product batch_sizes [1, 4, 8, 16] seq_lens [64, 256, 512] generate_len 32 records [] for bs, sl in product(batch_sizes, seq_lens): input_ids, attention_mask make_inputs(bs, sl) prefill_ms, decode_times benchmark_inference(model, input_ids, attention_mask, generate_len) avg_decode_ms sum(decode_times) / len(decode_times) total_tokens bs * (sl generate_len) total_time (prefill_ms sum(decode_times)) / 1000 # seconds throughput total_tokens / total_time records.append({ batch_size: bs, seq_len: sl, prefill_ms: round(prefill_ms, 3), decode_ms_per_token: round(avg_decode_ms, 3), throughput_tokens_per_sec: round(throughput, 2), }) df pd.DataFrame(records) df.to_csv(benchmark_results.csv, indexFalse) print(df)实际运行中batch size 16、序列长度 512 的配置在单张 A100 上 prefill 大概是 200-400ms 量级decode per token 可能在 10-30ms 之间取决于是否使用 KV cache。如果跑在消费级显卡如 3090 或 4090 上数字会略有不同但趋势一致prefill 随序列长度增长明显decode 随 batch size 增长明显。小技巧第一次跑完整矩阵前先单独跑一遍最大配置以估算总时长避免实验跑到一半才发现时间不够或者显存直接 OOM 中断整个循环。3.4 显存占用的额外监测除了时间指标显存占用也是 benchmark 的重要输出。PyTorch 提供两个关键接口torch.cuda.memory_allocated()返回实际分配的张量内存torch.cuda.max_memory_allocated()返回最近一次重置以来的峰值。在每次推理前先执行一次torch.cuda.reset_peak_memory_stats()跑完后再读取峰值这样拿到的是本次推理独占的显存增量而不是该进程的累计占用。torch.cuda.reset_peak_memory_stats() with torch.inference_mode(): out model(input_idsinput_ids) peak_mem torch.cuda.max_memory_allocated() / 1024**3 # GB显存数据可以放进同一个 DataFrame 中一并保存。这份数据在后续做 KV cache 优化时尤其重要——优化的目标通常就是在不显著增加显存的前提下提升生成速度。4. 常见问题与排查技巧实录4.1 结果波动大每次跑出来的数字都不一样这是 benchmark 新手遇到最多的问题。原因通常有三个一是没有做 warmupGPU 的缓存和 cuDNN autotune 还没热起来二是没有用 CUDA Event 计时混入了 CPU 调度时间三是在跑不同配置时GPU 有其他进程占用频率被拉低。解决办法正式计时前至少跑 3 次 warmup每次配置重复测 5 次取中位数或均值同时记录标准差。我自己的习惯是重复 5 次取中位数因为均值容易被偶发的调度抖动拉高。另外在跑 benchmark 的机器上最好关闭图形界面、浏览器等 GPU 渲染任务它们会干扰显存和算力共享。4.2 OOM 频繁出现OOM 最容易发生在 batch size 大且序列长度高的组合上。这里有一个朴素的避坑技巧在做 benchmark 之前先用一个简单的函数估算理论显存占用。一个标准做法是先用小配置测出每个 token 对应的激活显存占用再按线性外推估算大配置是否会 OOM。还有一种容易被忽略的情况加载模型、优化器状态、KV cache 的显存是分开计算的。如果接下来要对比不同 KV cache 实现那么你需要先把模型和优化器加载后的基线显存记录下来再单独量 KV cache 的内存增量。如果不把这两者分开你会误以为是 cache 实现导致显存暴涨实际可能只是 PyTorch 显存分配器的不确定性。4.3 用 model.generate() 测出的数字没法拆解 prefill 和 decode前面提到model.generate()是一个黑盒返回的是完整序列无法直接告诉你 prefill 用了多少时间、decode 每个 token 又用了多少时间。如果你发现自己在 benchmark 结果里只有一个“总耗时”说明你用错了工具。想要得到拆解的指标必须手动控制生成循环。这也意味着在接着说下一个主题之前建议把这套手动推理循环封装成函数后续所有性能对比都基于同一个入口函数。只有统一入口不同配置的对比才公平。4.4 注意 GPU 同步在哪里发生CUDA 的 kernel 执行是异步的调用model()返回后GPU 上的计算可能还没结束。所以计时事件必须用torch.cuda.synchronize()做同步。但要注意同步本身是有一点开销的。如果 decode 循环的每步都同步一次时间误差虽然不大但为了极致测量准确性可以改成整个 decode 循环开始前记录 start event结束时记录 end event最后除以生成 token 数得到平均每 token 延迟这样会略掉同步点之间的空档。但是这样的话就看不到每个 token 的耗时分布了。所以我通常的做法是既记录整体 decode 总耗时也额外做一次“每步记录耗时”的细粒度测试至少能观察是否有某个 token 的生成时间异常高通常是触发显存页迁移或 L2 cache 抖动。4.5 结果和别人的对不上不同型号 GPU、不同 PyTorch 版本、不同 CUDA 版本下跑结果有差异是正常的。但除了硬件和软件环境之外还有一个经常被忽略的因素torch.backends.cudnn.benchmark。这个选项默认是 False如果你设置了 TruePyTorch 会在第一次遇到某种 shape 时自动搜索最优算法首次调用可能非常慢但后续调用会变快且在不同 shape 之间切换时cudnn 会重新 autotune。做 benchmark 时建议一开始固定 cudnn 的 benchmark 为 False或者提前用实际输入 shape 做足 warmup尽量让运行环境保持稳定。5. 从 benchmark 数据能分析出什么跑完一组实验、拿到表格并不等于完成 Part1。真正的价值在于你能否从数据中讲出故事。我的实际数据以 A100 为例GPT-2 模型大致呈现出以下规律在 batch size 为 1 时decode per token 延迟受序列长度影响很小因为模型每步都只处理一个 token序列长度主要通过 KV cache 的历史参与注意力计算来影响性能。如果关闭 KV cache序列长度对 decode 的影响就会显著放大。当 batch size 增大时decode 延迟会上升但吞吐量也在上升。这背后是“延迟与吞吐的权衡”简单把 batch size 拉到最大并不意味着最优因为最终还要考虑显存限制和服务端响应时间。Prefill 时间是衡量首 token 延迟的关键指标。在对话类应用中prefill 时间直接决定了用户看到第一个字要多久。如果 prefill 时间占总耗时比例很高说明系统瓶颈在计算密集部分适合通过算子融合、量化等手段优化如果 decode 时间占比高则说明瓶颈在访存带宽适合通过减少 KV cache 访问次数、使用更好的缓存策略来优化。把这些观察写进课程报告就是一份合格的 Part1 输出。因为这些观察已经为后续的优化工作点明了方向哪个环节最值得投入精力哪种优化手段可能取得最大收益。6. 总结之外的个人实操建议最后分享几个我这次跑完 Part1 之后印象比较深的体会。第一个体会是评测代码值得花一周时间来打磨而不是“随便写写就行”。因为后续所有优化工作的判断都依赖这份评测代码的准确性。评测口径有误后面可能会得到错误的优化结论比如明明没有加速却因为计时方式有误以为快了。第二个体会是每次改动模型实现后都应该在相同输入下重新跑一遍 benchmark并且用 git 记录下代码版本。比如同一个函数可能在某个 commit 上解码延迟是 12ms在下一个 commit 变成 10ms但如果你没有版本记录和自动化的评测脚本你根本说不清性能变化是哪个改动引起的。养成配套的发布与测试习惯越早越好。第三个体会是不要迷信工具打印的数字要对自己设置计时点。比如我第一次跑 benchmark 时torch.cuda.Event 的 elapsed_time 返回单位是毫秒但我误当成秒来处理结果吞吐量算错了三个数量级浪费了一整晚。这种基础单位的错误非常低级却很常见。建议在脚本中强制做一次 sanity check比如实际生成 10 个 token如果单 token 解码时间小于 0.01ms或者大于 10s大概率是单位或者 GPU 使用出了问题。第四个体会涉及扩展方向。这次 benchmark 只覆盖了单模型、单 GPU 的基础推理完全可以扩展成多组对照比如 FP32 和 FP16 的对比、torch.compile 前后对比、不同 KV cache 实现方式的对比。这些扩展并不需要额外耗费太多精力只要把组织实验的基础架构搭好之后每加一组对比就只是修改配置字典的事情。课程作业的时间有限但把这份基准测试做得越扎实后续所有实验的产出效率就越高。在最后还是想多说一句benchmark 这门“手艺”看着简单真正做靠谱了需要对整个推理链路有全局理解。如果你也在写这类评测代码希望上面这些踩坑记录和数据整理方式能帮你少走几步弯路。