Reverse-engineering NVIDIA's cuda-checkpoint for faster cold starts | Doubleword

There’s a little known feature in the closed-source NVIDIA driver that lets you freeze a running CUDA process, serialize its GPU state into host memory, and later restore it to the GPU exactly as it was. We used it in an earlier post to speed up SGLang server startup by up to 70x. The utility is called cuda-checkpoint. The feature is documented, but how it works isn’t. One very frustrating aspect, that dogs anyone trying to use it to checkpoint complex GPU processes, is that the checkpoint transfers come nowhere close to saturating PCIe bandwidth. We left off our investigation in the earlier post without a good answer for why that was the caseIn the end, we just used cooperation from the application side to work around it.. With some tooling from our last post, we can find out why it costs so much, and how to make it faster without modifying the application, or the driver. How to checkpoint a CUDA process Here’s a small CUDA program: __device__ int counter = 100; __global__ void increment() { counter++; }

int main(void) { cudaFree(0); // force context creation 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>>>(); // one thread, 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); } } It binds a UDP socket, and every time a packet arrives it launches a one-thread kernel that increments a __device__ int, reads it back, and replies with the valueThis is NVIDIA's demo, shipped alongside the cuda-checkpoint tool. I've trimmed the error handling. Same 4090 and driver 590.48.01 as the kernel-launch post.. The counter lives in GPU memory and starts at 100. Ping it, and it says 101. $ ./counter & $ echo -n ping | nc -u -w1 127.0.0.1 10000 # send the packet, and view the response 101 We can freeze this process — copy its GPU state out, tear its CUDA context down to nothing, remove it from the GPU entirely — and then, some time later, bring it back exactly where it was: $ P=$(pgrep -xn counter) $ cuda-checkpoint --action checkpoint --pid $P # (lock first; see below) $ cuda-checkpoint --action restore --pid $P $ echo -n ping | nc -u -w1 127.0.0.1 10000 102 In between those two commands the process holds no GPU memory, has no CUDA context, and does not appear in nvidia-smi. The counter, which lived only on the device, survives anyway. This is the mechanism that a previous post leaned on to restore a 122B-parameter server in a few seconds — there, cuda-checkpoint was a black box called by CRIU. This post is about how we can find out what’s inside the box. Watching the process disappear Let’s watch the process disappear from the device.

cuda-checkpoint drives a small state machine over the target process; --action lock moves it from running to locked, and then --action checkpoint moves it from locked to checkpointed. cuda-checkpoint --action lock --pid $P cuda-checkpoint --action checkpoint --pid $P Watching the process across the two calls, with nvidia-smi, its /proc/$P/maps, its open file descriptors, and the RssAnon line of /proc/$P/status:

runninglockedcheckpointedRssAnon12,860 kB12,860 kB420,812 kBNVIDIA VMAs in /proc/$P/maps26260NVIDIA fds in /proc/$P/fd/24240visible in nvidia-smiyesyesno lock doesn’t change anything observable. But checkpoint does: every mapping of a /dev/nvidia* file is gone, every file descriptor pointing at the driver is closed, and the process vanishes from nvidia-smi. As far as the kernel driver is concerned, this process is no longer using a GPU. The resident anonymous memory jumped by 407,952 kB at the same moment, as the GPU state moves into ordinary host memory, into the process’s own address space. Can we find it? Poking the checkpoint The jump in RssAnon is almost exactly the size of one new anonymous mapping that appears in /proc/$P/maps at the checkpoint. If we attach strace to the target across the checkpoint we can catch it being allocated: $ 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 bytes is about 398 MiB, against the 388 MiB of device memory nvidia-smi had attributed to the process.

So the inference is the counter’s device footprint has been serialized into this buffer, plus about ten megabytes of something else. MAP_POPULATE asks the kernel to fault the whole thing in immediately rather than lazily, which matters later. What’s in there? We know what the increment kernel compiles to — cuobjdump -sass counter gives us its SASS — so we can prove that it’s in the mapping by taking the first few instruction words as a needle and searching the anonymous mapping for them. We can also find the counter. Nearby, in a page that is otherwise entirely zeros, is a single non-zero integer holding 0x67 — 103, because I’d pinged it a couple of times before checkpointing. With a few more of these kinds of tricks, we can see that the checkpoint buffer structure is pretty simple: The counter demo's checkpoint image, mapped by classifying every 4 KiB page and locating the driver's GPU-mapped surfaces inside it. The driver state is the constant ~10 MiB diff between the image and the nvidia-smi footprint. It seems to be GPU-mapped host memory (the increment() SASS is in there, as are kernel-launch parameter banks, and the channels' notifier pages). the counterdriver state · ~10 MiBzerosthe process's device allocations, roughly newest first398 MiBthe counterdevice allocations,roughly newest firstzerosdriver state~10 MiB398 MiB What’s more, the checkpoint image is plain anonymous memory in a process we own. Let’s mess with it. Rather than increment the counter through the GPU, we can reach into the frozen image and tweak it there. /proc/$P/mem lets us write the process’s memory directlyNeeds CAP_SYS_PTRACE or ptrace_scope=0., so we seek to the offset where we found the counter and write 424242.

Then: $ cuda-checkpoint --action restore --pid $P $ cuda-checkpoint --action unlock --pid $P $ echo -n ping | nc -u -w1 127.0.0.1 10000 424243 The restore uploads the number back onto the GPU, the next packet runs counter++ on it, and the process replies 424243, incrementing our injected value. So we know that the host-side anonymous buffer is the device memory across a checkpoint. But how did that memory get there? Who does the work cuda-checkpoint, the process we invoked, cannot read the target’s device memory. The tooling that does lives inside the target process and worse, inside the closed-source userspace driver. If you strace the utility, essentially everything it does to the target process is this: $ strace -f -e trace=openat,read,write cuda-checkpoint --action checkpoint --pid $P ... openat(AT_FDCWD, "/proc/$P/task", O_RDONLY|O_DIRECTORY) = 35 openat(AT_FDCWD, "/proc/$P/task/2863110/comm", O_RDONLY) = 36 read(36, "cuda00001400006\n", 1024) = 16 # the CUDA service thread openat(AT_FDCWD, "/proc/$P/fd/6", O_RDONLY) = 35 # its reply pipe openat(AT_FDCWD, "/proc/$P/fd/5", O_WRONLY) = 36 # its command pipe write(36, "\5\0\0\0", 4) = 4 # "where do I talk to you?"

read(35, "-\0\0\0\20\0\0\0", 8) = 8 # -> use fds 45 and 16 openat(AT_FDCWD, "/proc/$P/fd/45", O_RDONLY) = 37 openat(AT_FDCWD, "/proc/$P/fd/16", O_WRONLY) = 38 write(38, "\6\0\0\0\0\0\0\0"..., 2064) = 2064 # handshake read(37, "\0\0\0\0\1\0\0\0", 8) = 8 write(38, "\2\0\0\0\1\0\0\0"..., 2064) = 2064 # opcode 2, action 1: checkpoint read(37, "\0\0\0\0\2\0\0\0", 8) = 8 # status ... In words, it walks /proc/$P/fd/, finds a pipe that libcuda opened inside the target process when the CUDA context was first created, and writes a command word into it. The action lives in a single word, in exactly the order the CLI lists them:

word1action0lock1checkpoint2restore3unlock On the other end of that pipe, inside the target, a thread named cuda00001400006cuda-checkpoint --get-restore-tid returns a tid that matches this thread. has been sitting in poll() since the process started.

Its kernel stack, the entire time the process is running, shows it waiting for work: $ sudo cat /proc/$P/task/<tid>/stack [<0>] do_poll.constprop.0+0x315/0x3c0 [<0>] do_sys_poll+0x1ef/0x290 [<0>] __x64_sys_poll+0x4e/0x150 That thread does the entire checkpoint, from inside the process that is being checkpointed. What the driver sees If the service thread is doing the checkpoint, then it must be making driver calls. So let’s try to watch those. An LD_PRELOAD shim on the target that decodes each ioctl’s command number and parameter structThe shim wraps ioctl, matches NVIDIA's 'F' magic, and decodes the NVOS54 (RM_CONTROL) and NVOS21/NVOS64 (RM_ALLOC) parameter structs against the open kernel modules. The appendix has more details. gives us a per-phase histogram:

phaseRM_CONTROLRM_ALLOCother NV_ESC_*UVMtotal ioctlslock————2checkpoint133259756793restore2131241171081197unlock————2 There is no checkpoint ioctl. All the commands that the driver sees when a checkpoint is performed are ordinary resource-manager ioctls, the same ones that libcuda uses to create and destroy contexts, allocate and free memory, and so on. It’s Checkpoint and Restore in Userspace, but for GPU contexts. Coming back At checkpoint the context was torn down completely, so at restore time we start from host buffer and a process with no GPU context. Decoding the classes restore passes to NV_ESC_RM_ALLOC makes it clear that the result is basically just context creation (see the previous post for some of the details), run again from scratch, followed by refilling the fresh allocations from the host image.