逆向工程英伟达CUDA检查点加速冷启动

背景

NVIDIA 闭源驱动中隐藏着一个鲜为人知的功能：它可以冻结正在运行的 CUDA 进程，将 GPU 状态序列化到主机内存，之后再将状态精确恢复到 GPU 上。这个工具名为 cuda-checkpoint，虽然其功能有文档说明，但内部工作机制从未公开。使用该工具对复杂 GPU 进程做 checkpoint 时，一个令人头疼的问题是：checkpoint 数据传输远未达到 PCIe 带宽的饱和水平。此前一篇博客曾利用该工具将 SGLang 服务器启动速度提升 70 倍，但留下了性能瓶颈的疑问。本文通过逆向工程，揭示了 cuda-checkpoint 的内部原理，并探索了在不修改应用或驱动的前提下加速 checkpoint 的方法。

核心内容

一个简单的 CUDA 演示程序

以下是一个小型 CUDA 程序，它绑定一个 UDP 套接字，每当收到数据包时启动一个单线程内核，递增一个 __device__ int 计数器，读取并回复当前值：

__device__ int counter = 100;
__global__ void increment() { counter++; }
int main(void) {
    cudaFree(0); // 强制创建上下文
    int sock = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
    sockaddr_in addr = {AF_INET, htons(10000), inet_addr("127.0.0.1")};
    bind(sock, (sockaddr *)&addr, sizeof addr);
    while (true) {
        char buffer[16] = {0};
        sockaddr_in peer = {0}; socklen_t n = sizeof peer;
        recvfrom(sock, buffer, sizeof buffer, 0, (sockaddr *)&peer, &n);
        increment<<<1,1>>>(); // 一个线程，counter++
        int h = 0;
        cudaMemcpyFromSymbol(&h, counter, sizeof counter);
        size_t bytes = sprintf(buffer, "%d\n", h);
        sendto(sock, buffer, bytes, 0, (sockaddr *)&peer, n);
    }
}

计数器初始值 100，驻留在 GPU 内存中。运行后，发送 UDP 数据包即可得到递增后的值：

$ ./counter &
$ echo -n ping | nc -u -w1 127.0.0.1 10000
101

冻结与恢复进程

使用 cuda-checkpoint 可以冻结该进程：将 GPU 状态复制出来，拆除 CUDA 上下文，使进程完全从 GPU 中脱离，之后再精确恢复：

$ P=$(pgrep -xn counter)
$ cuda-checkpoint --action lock --pid $P
$ cuda-checkpoint --action checkpoint --pid $P
$ cuda-checkpoint --action restore --pid $P
$ echo -n ping | nc -u -w1 127.0.0.1 10000
102

在两次命令之间，进程不持有任何 GPU 内存，没有 CUDA 上下文，也从 nvidia-smi 中消失。但只存在于设备上的计数器依然存活。这正是此前博客中用于在数秒内恢复 122B 参数服务器的基础——当时 cuda-checkpoint 作为 CRIU 调用的黑盒。

观察进程消失的过程

cuda-checkpoint 驱动一个有限状态机：--action lock 将进程从 running 状态转为 locked，--action checkpoint 再从 locked 转为 checkpointed。通过 nvidia-smi、/proc/$P/maps、文件描述符和 /proc/$P/status 中的 RssAnon 观察：

• lock 不改变任何可观测内容。
• checkpoint 后：所有 /dev/nvidia* 映射消失，指向驱动的文件描述符关闭，进程从 nvidia-smi 中消失。同时，RssAnon 突增 407,952 kB，表明 GPU 状态已转移到主机内存中。

探查 checkpoint 映像

RssAnon 的增量恰好等于 /proc/$P/maps 中出现的一个新匿名映射的大小。通过 strace 捕获该映射的分配：

$ strace -f -p $P -e trace=mmap
...
mmap(NULL, 417739792, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS|MAP_POPULATE, -1, 0) = 0x...

大小为 417,739,792 字节（约 398 MiB），而之前 nvidia-smi 显示该进程占用 388 MiB 设备内存。因此推断：计数器的设备内存已经序列化到这个缓冲区中，外加约 10 MiB 的其他内容。MAP_POPULATE 标志要求内核立即触发缺页，而非惰性映射，这对后续性能有影响。

通过 cuobjdump -sass counter 得到 increment 内核的 SASS 指令，然后以指令字作为“针”在匿名映射中搜索，可以证明内核代码确实存在于其中。计数器同样被找到：在一个全零页面中，有一个非零整数 0x67（103，因为在 checkpoint 之前已 ping 过几次）。

通过更细致的分类（每 4 KiB 页面），可以发现 checkpoint 映像结构非常简单：

• 大部分是 GPU 映射的主机内存（包含 increment() 的 SASS、内核启动参数块、通道的通知页面等）。
• 驱动状态部分约占 10 MiB 的固定差异。

篡改 checkpoint 映像

由于 checkpoint 映像属于进程的普通匿名内存，我们可以直接通过 /proc/$P/mem 修改它（需要 CAP_SYS_PTRACE 或 ptrace_scope=0）。例如，找到计数器偏移量，写入 424242，然后恢复：

$ cuda-checkpoint --action restore --pid $P
$ cuda-checkpoint --action unlock --pid $P
$ echo -n ping | nc -u -w1 127.0.0.1 10000
424243

恢复将数值上传回 GPU，下一次数据包到达后执行 counter++，返回 424243。这证实了主机侧匿名缓冲区就是 checkpoint 期间的设备内存。

谁在执行工作？

cuda-checkpoint 工具本身无法直接读取目标进程的设备内存。实际工作由目标进程内部、闭源用户态驱动 libcuda 中的线程完成。通过 strace 观察 cuda-checkpoint 的行为：

$ strace -f -e trace=openat,read,write cuda-checkpoint --action checkpoint --pid $