在封闭源码的NVIDIA驱动中，有一个鲜为人知的功能，允许你冻结一个正在运行的CUDA进程，将其GPU状态序列化到主机内存中，然后在稍后恢复到GPU的状态，完全和之前一样。我们在之前的文章中使用这个功能将SGLang服务器的启动速度提高了70倍。这个工具称为cuda-checkpoint。功能有文档说明，但其具体工作原理并没有说明。一个非常令人沮丧的方面是，对于任何试图使用它来检查复杂GPU进程的人来说，检查点的传输速度远远没有达到PCIe带宽的饱和。在之前的文章中我们没有找到合理的解释。最终，我们只是依赖于应用程序方面的协作来规避这个问题。借助我们上一篇文章中的一些工具，我们可以找出为什么成本如此之高，以及如何在不修改应用程序或驱动的情况下加速它。 如何检查点一个CUDA进程 这是一个简单的CUDA程序： __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); } } 这段代码绑定了一个UDP套接字，每当数据包到达时，它就会启动一个单线程的内核，递增一个__device__ int，读取回来，并回复该值。这是NVIDIA的演示，与cuda-checkpoint工具一起发布。我省略了错误处理。与内核启动文章相同的4090和驱动590.48.01。 计数器存在于GPU内存中，起始值为100。ping它，它将返回101。 $ ./counter & $ echo -n ping | nc -u -w1 127.0.0.1 10000 # 发送数据包并查看响应 101 我们可以冻结这个进程——将其GPU状态复制出去，拆除其CUDA上下文，完全从GPU上移除——而在一段时间后，再精确恢复到原来的状态： $ P=$(pgrep -xn counter) $ 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将它从运行状态移动到锁定状态，然后--action checkpoint将其从锁定状态转变为检查点状态。 cuda-checkpoint --action lock --pid $P cuda-checkpoint --action checkpoint --pid $P 通过nvidia-smi、/proc/$P/maps、其打开的文件描述符和/proc/$P/status中的RssAnon行观察这两个调用的进程： 运行状态 锁定状态 检查点状态 RssAnon 12,860 kB 12,860 kB 420,812 kB NVIDIA虚拟内存映射在/proc/$P/maps中 26 26 0 /proc/$P/fd/中的NVIDIA文件描述符 24 24 0 在nvidia-smi中可见 是 是 否 锁定不会改变任何可观察的内容。但检查点会发生变化：每个指向/dev/nvidia*文件的映射消失，指向驱动程序的每个文件描述符关闭，该进程在nvidia-smi中消失。就内核驱动而言，该进程不再使用GPU。此时，常驻匿名内存增加了407,952 kB，因为GPU状态转移到普通主机内存中，进入该进程自己的地址空间。 我们能找到它吗？检测检查点 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。因此，推测是计数器的设备足迹已被序列化到这个缓冲区，加上约十兆字节的其他内容。MAP_POPULATE请求内核立即对整个映射进行故障处理，而不是惰性处理，这在后面很重要。里面有什么？我们知道 increment 内核编译成什么—— cuobjdump -sass counter给出其SASS——所以我们可以通过取第一个...