【阅读原文】戳：ACK AI Profiling：从黑箱到透明的问题剖析
1. 背景
Kubernetes 作为 AI 时代下主要的操作系统，承载了绝大多数的 LLM 训练和推理的业务负载，这些 LLM 负载的普及推动了对 AI 训练与推理的精细化性能检测与调优需求，那么如何对一个在线的 AI 业务负载进行精细化的 Profiling 从而达到线上问题发现性能调优的目的呢？
ACK AI Profiling 基于 eBPF 和动态进程注入，可以做到对 Kubernetes 中的 AI 应用进行用户无侵入、业务无感知、低 overhead 的性能分析，涵盖 Python 进程、CPU 调用、系统调用、CUDA 库和 CUDA 核函数五个方面的数据采集能力。通过对采集的结果数据进行分析，用户可以更好地定位应用容器的性能瓶颈并掌握对资源的利用程度，进而对应用进行问题定位和性能优化。
本文从一个通用的客户问题出发，描述了一个问题如何从前置排查到使用 AI Profiling 进行详细的排查，最后到问题定位与解决、业务执行过程的分析，从而展现一个从黑箱到透明的精细化的剖析过程。
2. 发现问题
一个使用 vLLM 部署大模型的推理服务，开启参数 --gpu-memory-utilization 0.95 在推理服务启动时提前占用了所使用 GPU 的 95% 的显存。在运行一段时间后发现该推理服务出现了显存 OOM 的情况。
在测试环境对这个推理服务进行观测，经过对该推理服务的多次压测后发现，在该推理服务开始接收推理请求时显存占用一直增加却不会下降，于是便有了 3 个问题：
增加的显存是谁在使用？
这个现象是否正常？
如果显存占用会一直增加，有什么办法可以释放显存？
3. 前置排查
3.1 排查环境
在测试环境中我们使用 QwQ-32B 来复现问题场景并进行排查。
模型：QwQ-32B
GPU：NVIDIA L20 * 2
开启了参数 --gpu-memory-utilization 0.95
（通过该参数的设置 vLLM 会预先占用 GPU 显存的指定比例，以确保处理请求时有足够的 KV Cache 缓存空间可用。这减少了动态分配显存的延迟，并避免因实时分配导致的内存碎片问题）
3.2 显存分配观测
首先复现场景并观测 GPU 显存占用，重新启动推理服务，并提交一次推理请求来模拟现场。启动 vLLM 推理服务的时候通过 nvidia-smi 观测显存占用如下，可以看到 vLLM 推理服务启动时会申请显存用作模型的加载和 KV Cache 的预留，每卡占用 43161 MiB 显存：
在推理请求发起后，再次通过 nvidia-smi 观测显存占用如下，两张卡的显存占用相比服务启动时分别增长 26 MiB 和 14 MiB：
接着观测基于 DCGM 的 GPU 监控，在监控大盘上可以明显看到在接收到请求后推理服务有一个比较明显的显存的增长，且显存增长值也可以对应上述 nvidia-smi 所观测到的现象。继续等待一段较长的时间，观测监控指标发现这部分占用增长的显存一直未被释放。
基于观测到的现象，这时候怀疑是否为 KV Cache 过满导致的显存申请。通过观测服务接收推理请求过程中 KV Cache 的利用率的日志如下图所示，利用率较低均在 10% 以内，所以原则上不存在 KV Cache 的显存不足导致的显存额外占用。
4. AI Profiling排查
4.1 问题定位与解决
参考 ACK AI Profiling 官网文档 AI Profiling 全流程使用指南 [1] 对正在运行的 vLLM 推理服务进行 AI Profiling，开启 Profiling 项为 Python、CPU 以及 CUDA Kernel，且使用 SysOM 进行结果展示与查看，这里对一次完整的推理请求进行 Profiling。
在整个推理过程中，观察 CUDA Kernel 的 TimeLine，可以分别找到每张卡上 cudaMalloc 操作的部分。查看详情找到具体显存申请的字节数，经过计算和从监控中观测到的显存增长数基本一致，由此可以确认这里即为触发显存增长的点。
使用标注线在 TimeLine 上标注所定位的 cudaMalloc 的时间部分，如下图所示：
根据标注线的位置，打开 Python 的 TimeLine 部分。在 TimeLine上Zoom In 到对应标注的部分，分别分析分布在两张卡上的两个进程。如下图所示：
GPU 0：进程 1
GPU 1：进程 2
可以比较清楚的看到 Python 调用栈的细节，在显存申请的时间段内，Python 执行的方法栈为：thread.run() -> llm_engine.step() -> worker.execute_model() -> model_runner.execute_model() -> decorate_context()，而且可以明显的看到 decorate_context 方法是 Pytorch 库中基础方法。进而在 Pytorch 源码中具体分析 decorate_context() 的细节代码，可以看到最终落到了 with ctx_factory() 的调用。
def decorate_context(*args, **kwargs): with ctx_factory(): return func(*args, **kwargs)
ctx_factory() 是 Pytorch 的上下文管理器，常用来管理 GPU 资源，性能监控，异常捕获等场景。因为 vLLM 的显存申请在推理服务启动过程中已经结束，这里额外申请的显存是在处理推理请求过程中触发申请的，所以基本可得出此现象与 Pytorch 有关。
结合 Pytorch 显存缓存的管理细节（此处不做赘述），可以得出和 Pytorch 显存缓存管理机制有关的结论：推理过程中增加的显存是由于 Pytorch 的显存缓存机制预留了一定数量显存。nvidia-smi 看到的显存占用值是 reserved_memory 和 torch context 显存之和。在推理过程中虽然不会主动释放，但是随着 Cache block 的增多，也会有大量新的请求使用已经存在的空闲 block，增长趋势会愈发变慢最终达到稳态。所以这个现象属于正常，需要给 Pytorch 留有一定的显存申请空间即可避免出现 CUDA OOM 的问题。
根据结论我们可以对应的给出一些使用建议来减小显存的需求或释放显存，例如：
避免显存碎片化：当出现显存剩余量足够，却无法分配连续的显存，可以通过环境变量 CUDA_PYTORCH_CUDA_ALLOC_CONF 将 max_split_size_mb 调小，来减少显存碎片化的影响；
尝试其他优化策略：除了调整 max_split_size_mb 外，还可以考虑其他优化策略来减少显存碎片化，如使用显存清理工具（如 torch.cuda.empty_cache() ）或调整模型和数据加载策略。
4.2 推理过程概览
参考以上的排查过程，还可以从 Profiling 结果中获取其他更多的信息。这里简单整体呈现一次推理过程的 Profiling 结果，以及表现出的一次推理过程大概的步骤。
查看收到推理请求后整体的 TimeLine 概览，从下面的 Python Profiling 的 TimeLine 中可以看到整个推理过程主要的步骤：输入预处理 -> 模型 forward 计算 step -> 结果生成与返回。详细的每一个方法的调用可以 Zoom In 到具体的地方进行细节的查看，这里不做额外的描述。
然后结合以下观测 CUDA Kernel 部分的 TimeLine，可以看到对应上面 Python 调用过成的预处理部分以及后续的每一次 vLLM 的计算 step，其中包含着具体的 gemm 的运算过程以及数据结果返回的 cuMemcpyDtoH 显存传输调用。又因为是两卡部署的推理服务，所以在 step 间还充斥着 GPU 间使用 nccl 做数据传输的具体操作（broadcast、allreduce、sendRecv 等 Kernel）。
在每一个计算 Step 结束后会发现有一段较小的时间空档，这段时间内 GPU 上没有任何的 Kernel 动作在执行，Zoom In 到细节处可以发现这段时间主要包括 HTTP 请求的返回和服务 Metrics 的计算与上报，如下图所示。
最后观察 GPU 核函数计算统计图表，可以观察出调用主要集中前 3 个方法：
std::enable_if::type internal::gemvx::kernel
vllm::paged_attention_v2_kernel
ncclDevKernel_AllReduce_Sum_bf16_RING_LL
可以看出大部分的 GPU 计算时间都在做 gemvx 的矩阵运算，是为具体的推理运算过程。其中夹杂着一些 vLLM 框架中实现的 paged_attention 方法，以及多卡部署推理服务时所进行的 ncclDevKernel_AllReduce 通信操作。通过该统计图表即可很明显的分析自己 GPU 的繁忙程度以及所完成的具体操作的时间占比，进而分析出可能出现的性能瓶颈，例如假设 ncclDevKernel 占用了过多的 GPU 时长，即可以大致判断 GPU 通信为可以优化的部分，其余结论以此类推。
5. 总结
ACK AI Profiling 能力为精细化分析在线应用提供了更多的可能和选择，通过采集到的丰富的指标数据结合自己的业务详情可以详细的进行出现问题的定位与性能的分析，从而为 AI Infra 的可观测带来更大的价值。
参考文档：
我们是阿里巴巴云计算和大数据技术幕后的核心技术输出者。
获取关于我们的更多信息~