nvsnap: inject cooperative checkpoint/restore hooks into inference engines (vLLM/SGLang/NIM)

Summary

Inject cooperative checkpoint/restore lifecycle hooks into stock inference
engines (vLLM, SGLang, and vLLM/SGLang-backed NIM) through nvsnap's existing
transparent injection stack (worker-extension / sitecustomize staged by an
init container + mutating webhook), so multi-GPU checkpoint/restore works
without forking the engine and without waiting for upstream hooks to merge.

Motivation

The multi-GPU checkpoint blocker is cross-GPU collective peer state: with
tensor parallelism, the FlashInfer all-reduce workspace, all-to-all managers,
and NCCL communicators hold cross-GPU mappings, and cuda-checkpoint cannot
serialize them (it hangs while any graph-captured cross-context collective is
live). The fix the ecosystem is converging on is cooperative teardown: the
engine releases those collective resources on a checkpoint hook and rebuilds
them on restore.

Upstream vLLM is building exactly this (checkpoint_prepare /
checkpoint_restore on device communicators, vllm-project/vllm#46877), but it
is unmerged, has no GPU validation yet, and is blocked on FlashInfer finalizing
its own checkpoint APIs. Upstream adoption across vLLM, SGLang, and NIM will
take a long time.

We can deliver the same hooks now by injecting them, the same way nvsnap
already injects its worker extension. This is a monkeypatch layer over stock
engine internals, not a fork.

Scope

• vLLM: inject via --worker-extension-cls + sitecustomize; reach the device
  communicator and free the FlashInfer all-reduce workspace + all-to-all
  managers by identity, then re-establish on restore.
• SGLang: sitecustomize monkeypatch of the TP worker (same teardown targets,
  different class names).
• NIM (vLLM/SGLang backend): inject into the underlying engine.
• Out of scope: TRT-LLM-backed NIM (collective teardown is in compiled code,
  not reachable from injected Python) and forking any engine.

Approach

Extend the injected worker extension with the upstream hook's teardown target
list plus the graph handling nvsnap already does:

1. checkpoint_prepare: drop captured CUDA graphs, free the FlashInfer
   all-reduce workspace + all-to-all managers, tear down the process group.
2. Checkpoint via the existing agent flow.
3. checkpoint_restore: re-init the process group, re-establish workspaces,
   recapture graphs (background recapture so restore latency is unaffected).

Where the engine does not expose a teardown method, free the workspace's
underlying device memory directly from the held object handle.

Open question this also answers

This injection is both the delivery mechanism and the decisive test of whether
cooperative FlashInfer-workspace teardown clears the multi-GPU cuda-checkpoint
hang. A weaker teardown (graph drop + process-group destroy, without freeing the
FlashInfer workspace) still hangs; freeing the workspace is the ingredient the
upstream hook identifies as the missing piece. Validating it against the
multi-GPU checkpoint repro is a concrete outcome (and the GPU validation the
upstream PR is missing).

Caveats

• Monkeypatching engine internals is version-fragile; pin the engine version
  per workload.
• Runtime is untouched (hooks run only at the checkpoint boundary), so there is
  no steady-state performance impact.

References

• vllm-project/vllm#46877 (upstream cooperative checkpoint hooks for
  communicators, starting with FlashInfer)

First step

Extend the injected worker extension with the FlashInfer workspace + all-to-all
teardown, alongside the existing graph drop, and validate against the
multi-GPU checkpoint repro on H100.