fix(agent): restore access permissions and unblock collective multicast calls

[galletas1712]

Summary

Two independent defects in the cuinterpose shim, both found while bringing up
GLM-5.2-NVFP4 TEP8 (vLLM, TP=8, 8x B200) with GMS V1 + CRIU checkpoint/restore.
Each one on its own made that configuration unusable. Both are fixed here, on
top of #79, as two separate commits.

Stacked on #79 (feat/snapshot-cuda-vmm-multicast-posix), which is itself on #78.

---

1. CUDA_ERROR_ILLEGAL_ADDRESS on the first request after restore

Commit: fix(agent): restore CUDA access permissions for all covered mappings

Symptom

A restored engine came up healthy — CRIU restore clean, wake_up() in 0.9 s,
endpoints registered — and then every one of the 8 ranks died with
CUDA_ERROR_ILLEGAL_ADDRESS on the first inference request. The same model and
the same engine flags, cold-started without checkpoint/restore, served
correctly. So the fault was introduced by checkpoint/restore.

Root cause

cuMemSetAccess recorded its descriptors only when the requested range matched a
tracked mapping's (address, size) pair exactly:

mapping = find_mapping(address, size);   /* requires mapping->address == address
                                            && mapping->size == size */
if (mapping == NULL) {
    ...
    return function(address, size, descriptors, count);   /* passthrough, nothing recorded */
}

Any caller that maps at one granularity and sets access at another — which is
legal CUDA, and which PyTorch's caching allocator does routinely — took that
passthrough and left access_count at 0.

restore_mappings then replays access only when access_count != 0:

result = map(mapping->address, mapping->size, mapping->offset, handle, 0);
mapping->mapped = true;
if (mapping->access_count != 0)          /* false => access never replayed */
    result = set_access(...);

So those mappings were restored mapped but with no device access. The VA
resolves, nothing looks wrong, and the first kernel to dereference it faults.

Evidence

The cuinterposer.state sidecar records access_count per mapping, so this is
measurable directly from a checkpoint, before any restore. Decoding the sidecar
of the failing 8-rank GLM-5.2 checkpoint:

	mappings
total	336
with access	144
with access_count == 0	192

The 192 broken ones were exactly 64 each at 2 MiB, 8 MiB and 24 MiB — 8 buffers
per rank across 8 ranks, i.e. the non-GMS scratch allocations. GMS's own weight
mappings were never affected, because GMS maps and sets access with an identical
(base, aligned_size) and so always hit the exact-match path.

Fix

Record access against every tracked mapping the range fully covers, and fail
closed with CUDA_ERROR_NOT_SUPPORTED on a range that only partially overlaps
a tracked mapping, rather than passing it through and silently losing the access
across restore. Classification runs before the driver call, so an unrepresentable
range is rejected without mutating device state, and a range that covers nothing
we track stays a pure passthrough.

Result

Same workload, same record counts, before and after the fix:

	before	after
mappings with access	144	336
mappings without access	192	0

Restore then serves correctly: finish_reason: stop, content "2, 3, 5, 7, 11",
and three follow-up requests all returning the right answer, with zero CUDA
errors and no container restarts.

---

2. Rank-divergent hang during multi-rank startup

Commit: fix(agent): do not hold state_lock across collective multicast calls

Symptom

An 8-rank startup wedged for 15+ minutes inside torch.compile. py-spy showed
the tensor-parallel group split across two different code paths in vLLM's
get_fi_ar_workspace: 6 ranks still attempting the mnnvl allreduce workspace,
2 ranks already fallen back to trtllm. The fallback path runs its own
collective, which the other 6 never joined, so the group deadlocked.

Root cause

cuMulticastCreate and cuMulticastAddDevice ran the real driver call while
holding the process-global state_lock. Both contend the NVIDIA RM top lock and
are NVSwitch-team collectives, so under concurrent multi-rank startup they can
block for a long time.

serve() needs that same lock to answer a peer's export request, and the accept
loop served connections inline. So a creator stuck inside one of those calls
could not respond, the peer hit the hard-coded 30 s request_export deadline,
and got CUDA_ERROR_INVALID_HANDLE — while other ranks, which happened to ask
earlier, succeeded.

A per-rank, timing-dependent failure is pathological for a collective framework:
which ranks lose is pure timing, and any framework that treats workspace creation
failure as a rank-local decision will diverge and hang.

Note the codebase already knew this shape — bind_memory/bind_address drop the
lock around their driver call, with the comment "Bind waits for the complete
multicast team, so it must not block the creator endpoint under state_lock."
The same treatment was simply never applied to create and add_device.

Fix

• Drop state_lock around the cuMulticastCreate / cuMulticastAddDevice
  driver calls, mirroring bind_memory. Revalidate the phase on reacquire (and,
  for add_device, re-look up the handle, which may have been released while
  unlocked) via a new state_is_active callback.
• Make the export and control deadlines configurable with
  DYN_SNAPSHOT_EXPORT_TIMEOUT_SECONDS and
  DYN_SNAPSHOT_CONTROL_TIMEOUT_SECONDS, and raise the default from 30 s to
  300 s. The unicast import path had the same exposure.
• Serve each control connection on a detached thread so a slow request no longer
  blocks accept(). serve() still takes state_lock for its whole body, so
  requests stay serialised — the win is that peers are no longer stuck in
  connect() burning their deadline before they are even heard.

---

Note for reviewers

The immediate trigger for issue 2 is arguably a vLLM bug: its mnnvl -> trtllm
fallback is a rank-local decision with a collective inside it and no agreement
step, so any transient per-rank failure deadlocks the group. FlashInfer ships
the correct primitive for this (all_ranks_support_mnnvl, which AND-reduces the
per-rank probe and documents exactly this hazard), but vLLM bypasses it by
pre-resolving "auto" locally. That is worth reporting upstream separately.

It is still right to fix the shim: it should not manufacture per-rank,
timing-dependent CUDA failures during a collective startup.

Worth a separate discussion: #79 also made cuMulticastGetGranularity forward to
the driver, where #78 returned CUDA_ERROR_NOT_SUPPORTED. That is the capability
probe, so #79 advertises multicast support and opens the whole multicast path.
A uniform, early, deterministic capability denial is much easier for frameworks
to handle than a late per-rank timeout. Gating multicast behind an env switch
(defaulting to the #78 stub) may be worth considering — at the cost of losing
NVSwitch multicast allreduce, and requiring NCCL_NVLS_ENABLE=0 and
TORCH_SYMM_MEM_DISABLE_MULTICAST=1 to degrade gracefully rather than fail hard.

Validation

• Both commits build independently and cleanly under -Wall -Wextra -Werror
  (the stack is bisect-clean).

• End-to-end on four models, all vLLM TP=8 / TEP8 on 8x B200 with GMS V1 +
  CRIU. Every one captures cleanly with no allreduce deadlock, restores with 0
  restarts, and serves correct inference after restore — the check that
  actually matters here, since the illegal-address fault only surfaced on the
  first request:

  Model	Restore total	criu_restore	cuda_restore	wake_up	Post-restore inference
  Qwen3.8-27B	27.7 s	12.1 s	15.3 s	0.481 s	correct
  DSV4-Flash-NVFP4	32.7 s	12.5 s	19.8 s	0.704 s	correct
  GLM-5.2-NVFP4	36.5 s	12.5 s	23.6 s	0.912 s	correct
  DSV4-Pro-NVFP4	57.0 s	31.0 s	19.8 s	0.751 s	correct

  Inference check was a deterministic prompt ("2, 3, 5, 7, 11", and 3x
  "6 times 7" -> 42 on GLM-5.2), with zero CUDA errors in the engine logs.
  Before the fix, GLM-5.2 died with CUDA_ERROR_ILLEGAL_ADDRESS on all 8 ranks
  at exactly this point.

• Sidecar verification before/after, decoded from cuinterposer.state:
  GLM-5.2 192 mappings without access -> 0 (336/336 now carry access);
  Qwen3.8 336/336; DSV4-Pro 144/144.

• Shadow-engine promotion (kill the active engine, promote the paused restored
  shadow in the same pod) also works on the patched shim: verified on
  DSV4-Flash, 21 consecutive correct responses on the promoted engine. It
  exposes a separate, unrelated defect in GMS V1 lock admission — the
  reclamation of a SIGKILLed peer's epoch blocks new-writer admission and is
  O(GPU bytes held) — which is being fixed independently in the dynamo repo and
  is not a cuinterpose issue.

Follow-ups not addressed here

• The shim ships in the runtime image; the validation run loaded the patched
  build via LD_PRELOAD from a shared volume. The image needs rebuilding from
  these commits.
• Test coverage is 2 tests at WORLD_SIZE=2, which cannot see either of these
  bugs (one needs a non-exact cuMemSetAccess range, the other an 8-rank timing
  race). Regression tests for both would be valuable.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (1)

• main

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: 97c4c57b-5e55-4788-98b5-974652ce866f

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}