探How 社区

作者在B300上对cuda-checkpoint实现进行优化，最终将GPU检查点速度提升至15倍。工作从理解cuda-checkpoint的工作原理和瓶颈入手，尝试了多种提速思路并记录了成败。文中提到运行该版本cuda-checkpoint可观察到约2倍加速。

以下正文由 AI 翻译自 Hacker News 最新（原文版权归原站所有），可能存在偏差，重要内容请以原文为准。

GPU checkpointing 让应用可以保存 GPU 状态，之后从暂停处继续运行。这让我们在 GPU 上运行应用时拥有了更多灵活性。推理服务可以在空闲时段释放 GPU，等流量回来后再恢复（Modal: Serverless GPU scaling；Doubleword: Faster SGLang cold starts；Cedana: GPU checkpointing and migration with CRIU）。

出于好奇，也为了顺便学点东西，我们想看看能把 GPU checkpointing 和恢复推进到什么程度。我们从 B300 上的 cuda-checkpoint 实现入手，花了一些时间弄懂它的工作原理，以及瓶颈在哪里。然后我们尝试了各种思路来提速，看看哪些有效、哪些无效。

cuda-checkpoint 通常会与另一个程序配合使用，比如 CRIU（CRIU: GPU checkpointing）或 gVisor（gVisor: GPU checkpoint/restore；Modal: Memory snapshots with gVisor；Modal: GPU memory snapshots），以便把 GPU 资源和应用本身完整序列化到磁盘。这类应用常常会提到沙箱或快照。

我们观察到，运行这个版本的 cuda-checkpoint 有可能获得 2 倍加速。

对 CUDA 进程做 checkpoint 和 restore 各分两步：

做 checkpoint 时，我们首先锁定进程的 CUDA 活动，阻止新任务提交，并等待已有任务完成。然后将其 GPU 状态 checkpoint 到主机内存，并释放其 GPU 资源（虽然释放 GPU 资源对 checkpoint 本身并非必需，但这样可以让其他资源紧张的应用使用 GPU）。

在此期间，另一个 GPU 应用可以使用被释放出来的资源。

要恢复时，我们先恢复 GPU 资源，再把保存的数据拷回 GPU 内存。最后解锁进程的 CUDA 活动，让它能重新提交任务。

NVIDIA 的 cuda-checkpoint 工具通过简单的命令行接口暴露这四项操作（该 CLI 通过管道使用自定义协议与 CUDA 进程通信。它只是一个薄封装，但会影响延迟，后文会展示）：

ps aux | grep (your gpu application)
cuda-checkpoint --action lock --pid $PID
cuda-checkpoint --action checkpoint --pid $PID
cuda-checkpoint --action restore --pid $PID
cuda-checkpoint --action unlock --pid $PID

实际交互方式如下：

pid_t pid = getpid();
CUcheckpointLockArgs lock_args{0};
cuCheckpointProcessLock(pid, &lock_args);
CUcheckpointCheckpointArgs ckpt_args{0};
cuCheckpointProcessCheckpoint(pid, &ckpt_args);
CUcheckpointRestoreArgs rest_args{0};
cuCheckpointProcessRestore(pid, &rest_args);
CUcheckpointUnlockArgs unlock_args{0};
cuCheckpointProcessUnlock(pid, &unlock_args);
CUprocessState st;
cuCheckpointProcessGetState(pid, &st);
if (st == CU_PROCESS_STATE_CHECKPOINTED) {
 // ...

该工具只把 checkpoint 镜像存到 CPU 内存。

Doubleword 的 Fergus Finn 在 RTX 4090 上把约 9 GB 的 checkpoint-and-restore 基准测试提速了近 4 倍，从 4.5 秒降到 1.2 秒。这些改进来自：避免重复的 CLI 初始化、启用 huge pages、预分配主机缓冲区，以及推迟缓冲区清理（Reverse-engineering NVIDIA's cuda-checkpoint for faster cold starts）。我们会更详细地探讨这些技术，追踪驱动内部实现，解释它们为什么有效、解决了哪些瓶颈。我们也会自己应用其中一些，并验证 CLI、huge pages 和预分配主机缓冲区这几项优化确实有用。

我们会深入 libcuda（NVIDIA 的用户态 CUDA 驱动），看看每一步底层到底发生了什么。我们想通过 trace 观察时间花在哪里（每个 perfetto trace 显示一个 vLLM 实例运行在两块 B300 上，跑 GLM5-3 flash，总计约 347GB。每个 GPU 实例占用约一半内存（约 172GB），我们对两块 GPU 都做 checkpoint，每块称为一个 rank）。

注意：每个优化方案都分别以未 trace / 已 trace 的方式运行过，开启 tracing 带来的差异不大。

我们在同一台 B300 主机的 GPU 0 和 1 上，用 vLLM 0.30.0 运行 GLM-5.3-flash。张量并行度为 2，所以每块 GPU 有一个 worker 进程，称为一个 rank。我们使用 eager 执行、8,192 token 的上下文上限，每块 GPU 8 GiB KV cache。

CUDA_VISIBLE_DEVICES=0,1 python -m vllm.entrypoints.openai.api_server \
 --model /path/to/GLM-5.3-flash \
 --tensor-parallel-size 2 \
 --enforce-eager \
 --max-model-len 8192 \
 --kv-cache-memory-bytes 8589934592

生成参数放在每个请求里：

{"temperature": 0, "seed": 1234, "max_tokens": 128}

我们先发一些请求确认 vLLM 在运行。基准控制器会等到 vLLM 报告没有运行中或排队中的请求。然后控制器对两个 rank 并行启动 cuda-checkpoint --action lock --pid <worker-pid>，并等待两个调用都完成。它按同样方式依次执行 checkpoint、restore 和 unlock。每个循环结束后，它发送同样的三个请求，并检查其答案和 token ID 与 checkpoint 之前记录的一致。

我们的 tracing 工具使用 Frida，在两个 vLLM worker 内挂上函数入口和返回钩子。它也对每个 cuda-checkpoint 进程做插桩，这样我们既能看到 CLI 的初始化和请求，也能看到目标 worker 内部的工作。钩子记录驱动操作和 CPU 调用，比如 mmap 和 memcpy。我们还加入来自 NVML 的 GPU 计数器，采样目标间隔为 100 ms，以及来自 /proc 的 CPU、内存和缺页计数器。

这些 trace 和计数器会汇总成单个 perfetto trace 输出。

0% ↓ ⋯ Always load traces here automatically Loading preview…

我们可以看到锁定、checkpoint、restore 和解锁各花了多少时间。显然，checkpoint 和 restore 占了大部分时间，所以我们先深入看 checkpoint。

checkpoint 的 profile 有三个主要阶段：setup、copying 和 cleanup。Copying 占的时间最多。即便按这台机器理论上的 PCIe 最大带宽 64GB/s 计算，checkpoint 172GB 也至少需要约 2.5 秒。

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）

正文由探How 从来源站点 RSS 同步，内容未经改写；遇到排版缺失或需要图片、视频时请以原文为准。

在B300上将CUDA检查点速度提升15倍的工作记录