feat(comm): preserve all-reduce graph VAs across checkpoint restore by galletas1712 · Pull Request #3745 · flashinfer-ai/flashinfer

galletas1712 · 2026-06-26T23:05:24Z

Summary

Adds checkpoint lifecycle support for all-reduce workspaces whose pointers are captured by CUDA graphs. Physical CUDA VMM backing can be released before checkpoint and recreated at the same reserved virtual addresses after restore.

The workspace API is:

checkpoint_prepare() collectively quiesces the group and detaches physical backing while preserving graph-visible virtual addresses.
checkpoint_restore(comm_backend) attaches fresh backing at those addresses, reinitializes protocol state, and collectively fences before returning.
Repeating a successfully completed prepare or restore is a no-op.
Lifecycle failures propagate without rollback or retry; the affected workspace must not be reused.
Workspaces backed by torch symmetric memory explicitly reject this lifecycle because that allocator does not provide the required stable-address remap contract.

Implementation

Decomposes SymmDeviceMemory construction around _create_and_map_handles() and adds the inverse _unmap_and_release_handles() operation. Initial allocation and restore therefore use the same mapping path.
Retains CUDA virtual-address reservations while releasing unicast and multicast physical handles.
Uses a fresh communicator and handle exchanger during restore.
Supports unicast-only TRT-LLM workspaces without requiring multicast capability.
Moves the MNNVL all-reduce workspace to the same VMM-backed memory core.
Reinitializes TRT-LLM Lamport/barrier/control state and MNNVL Lamport/flag state before graph replay.
Rejects all-reduce launches while native backing is detached.

Validation

ruff format leaves the changed Python files unchanged.
ruff check passes for the changed Python files.
py_compile passes for the changed Python files.
git diff --check origin/main passes.
The torch-symmetric-memory rejection test passes for both all-reduce workspace types.
Checkpoint and MNNVL all-reduce test collection succeeds (278 tests).

Summary by CodeRabbit

New Features
- Added checkpoint_prepare() / checkpoint_restore(comm_backend) lifecycle support for symmetric-memory–backed all-reduce fusion workspaces (TRT-LLM and MNNVL), enabling safe reattachment and protocol reinitialization.
Bug Fixes
- Tightened validation for consistent handle attachment/detachment; fused all-reduce now rejects launches when required workspace handles are not attached.
Documentation
- Documented required call order, rank synchronization expectations, idempotency, and unsupported backings.
Tests
- Added distributed CUDA-graph coverage to verify behavior before/after symmetric-memory remapping.

coderabbitai · 2026-06-26T23:05:32Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: d0f4b667-fc15-48cd-881c-6fd77c006d8d

📥 Commits

Reviewing files that changed from the base of the PR and between cdd7e18 and 2dc23d5.

📒 Files selected for processing (1)

flashinfer/comm/trtllm_ar.py

🚧 Files skipped from review as they are similar to previous changes (1)

flashinfer/comm/trtllm_ar.py

📝 Walkthrough

Walkthrough

The PR adds checkpoint/restore lifecycle support for symmetric-memory-backed TRTLLM and MNNVL all-reduce workspaces, updates SymmDeviceMemory mapping state handling, and adjusts docs, tests, and legacy pointer wiring to match the new handle-based flow.

Changes

Workspace memory lifecycle

Layer / File(s)	Summary
SymmDeviceMemory lifecycle `flashinfer/comm/mnnvl.py`	`SymmDeviceMemory` tracks mapped state, conditionally validates multicast support, and centralizes handle creation and release around a mapped flag.

TRTLLM symmetric-memory workspace

Layer / File(s)	Summary
TRTLLM symmetric-memory workspace `flashinfer/comm/trtllm_ar.py`, `flashinfer/comm/allreduce.py`	TRTLLM all-reduce fusion adds symmetric-device-memory handle allocation, protocol initialization, metadata tracking, checkpoint restore, and a mapped-state guard before kernel dispatch.

MNNVL workspace checkpointing

Layer / File(s)	Summary
MNNVL workspace checkpointing `flashinfer/comm/trtllm_mnnvl_ar.py`, `tests/comm/test_trtllm_mnnvl_allreduce.py`	MNNVL all-reduce fusion switches to `McastGPUBuffer`-backed pointers, adds checkpoint prepare/restore APIs, requires attached handles before kernel launches, and updates the legacy test pointer source.

Docs and checkpoint coverage

Layer / File(s)	Summary
Docs and checkpoint coverage `docs/api/comm.rst`, `tests/comm/test_trtllm_allreduce_checkpoint.py`	The API docs describe checkpoint/restore semantics for both workspaces, and new tests cover destroyed-state errors plus distributed checkpoint and replay behavior.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related issues

[RFC] Checkpoint lifecycle APIs for CUDA-graph-stable communication workspaces #3791: Describes the same checkpoint/restore lifecycle APIs and stable-VA workspace remapping implemented here.

Possibly related PRs

flashinfer-ai/flashinfer#1991: Both PRs modify flashinfer/comm/trtllm_ar.py around fusion workspace metadata and protocol reinitialization fields.
flashinfer-ai/flashinfer#2056: Both PRs touch flashinfer/comm/trtllm_mnnvl_ar.py and its CommBackend-based handle transfer/synchronization flow.
flashinfer-ai/flashinfer#2118: Both PRs change the MNNVL handle lifecycle and IPC/multicast plumbing used by these workspaces.

Suggested reviewers: yzh119, jimmyzho, nv-yunzheq, bkryu, saltyminty

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately summarizes the main checkpoint-restore virtual-address preservation change.
Description check	✅ Passed	The description covers the required summary, implementation, and validation sections; only optional template sections are omitted.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

galletas1712 · 2026-06-26T23:06:23Z

@coderabbitai review

coderabbitai · 2026-06-26T23:06:29Z

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

gemini-code-assist

Code Review

This pull request introduces mechanisms to manage, validate, detach, and reattach graph-visible symmetric-memory workspace addresses (virtual addresses and layout states) captured by CUDA graphs across several communication modules, enabling stable virtual addresses during checkpointing or communicator refreshes. The review feedback identifies two critical issues: a runtime TypeError in allreduce.py caused by passing a standard Python bytes object directly to ctypes.cast, and a potential TypeError in trtllm_ar.py if device.index is None when initializing SymmDeviceMemory.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

coderabbitai

Actionable comments posted: 7

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@flashinfer/comm/allreduce.py`:
- Around line 183-244: The new public checkpoint APIs in the Comm backend are
missing the required API decorator. Add `@flashinfer_api` to
get_graph_visible_addresses, validate_graph_visible_addresses, detach_handles,
and reattach_handles in the allreduce module, and do not pass trace= since these
methods return dict/None and are not tensor I/O. If the decorator is not already
imported in this module, wire it in alongside the other API decorators so these
public methods remain crash-safe and consistently logged.
- Around line 218-244: The restore flow in reattach_handles should validate the
restored workspace before using the raw flag pointer and should fence peer
reinitialization before resuming execution. Reorder the logic so
validate_graph_visible_addresses runs before reading flag_ptr from
workspace_tensor or performing the cudaMemset/cudaMemcpy writes, then add a
collective synchronization/barrier after the Lamport and flag reset sequence so
all ranks finish reattaching before any can proceed.

In `@flashinfer/comm/mnnvl.py`:
- Around line 1249-1255: The reattach flow in the unicast/multicast setup leaves
partially created CUDA VMM handles live if an exception is raised before _mapped
is set. Update the reattach sequence around _create_and_import_unicast_handles,
_map_unicast_buffers, _create_and_import_multicast_handle, and
_map_and_bind_multicast_buffer so any failure triggers cleanup of already
created mappings/handles before re-raising, ensuring __del__ does not skip
release because _mapped is still false.
- Around line 1050-1256: The new public checkpoint/metadata methods in mnnvl.py
are missing the required `@flashinfer_api` decoration. Add `@flashinfer_api` to the
public APIs such as get_graph_visible_addresses,
validate_graph_visible_addresses, detach_handles, and reattach_handles, and do
not pass trace= because these methods manage handles/metadata rather than tensor
I/O. Use the method names above to locate the affected entry points and keep the
decorator consistent with the rest of flashinfer/**/*.py.

In `@flashinfer/comm/trtllm_ar.py`:
- Around line 441-447: The TRTLLM restore path is still constructing
SymmDeviceMemory in a way that forces multicast support even though it only
needs unicast buffer pointers from get_buffer_ptrs_host(). Update the
SymmDeviceMemory call in the restore flow to pass enable_multicast=False, and
adjust SymmDeviceMemory so the CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED check is
only enforced when multicast is actually enabled.

In `@flashinfer/comm/trtllm_mnnvl_ar.py`:
- Around line 246-301: The new public checkpoint APIs in trtllm_mnnvl_ar should
be decorated so they participate in the API inventory and logging. Add
`@flashinfer_api` to get_graph_visible_addresses,
validate_graph_visible_addresses, detach_handles, and reattach_handles, and do
not pass trace= since these methods return metadata or mutate handle state
rather than performing tensor I/O. Use the existing class/method names to locate
the implementations and keep the behavior unchanged.
- Around line 146-152: Validate the communicator used by McastGPUBuffer before
it exchanges MNNVL handles: the current TRTTLLM MNNVL path passes
mapping.tp_size and mapping.tp_rank while relying on comm_backend, which can be
the default MPIBackend() COMM_WORLD and may not match the TP group. Update the
constructor flow around the McastGPUBuffer call in trtllm_mnnvl_ar.py to assert
or verify that comm_backend corresponds to the same rank set as
mapping.tp_size/mapping.tp_rank before allocation, and fail fast or select the
correct communicator before any handle exchange happens.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: b892e5f2-c747-4c59-ace1-e6114cb51a62

📥 Commits

Reviewing files that changed from the base of the PR and between a7f5c5a and f9005d4fd15592fd334a06583b5fea438255aee4.

📒 Files selected for processing (4)

flashinfer/comm/allreduce.py
flashinfer/comm/mnnvl.py
flashinfer/comm/trtllm_ar.py
flashinfer/comm/trtllm_mnnvl_ar.py

Co-authored-by: William Arnold <7565007+Aphoh@users.noreply.github.com> Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/comm/test_trtllm_allreduce_checkpoint.py`:
- Around line 43-48: The process group setup in dist.init_process_group should
fail fast instead of waiting on PyTorch’s default timeout. Update the test’s
init call in test_trtllm_allreduce_checkpoint to pass a short timeout, using a
datetime.timedelta of 60 seconds alongside the existing backend, init_method,
rank, and world_size arguments.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 430598ec-138a-4c90-9b58-4cb392502878

📥 Commits

Reviewing files that changed from the base of the PR and between ddf02366b80f8e3d2c1a6f4f15928493df1a93ba and cdd7e18.

📒 Files selected for processing (7)

docs/api/comm.rst
flashinfer/comm/allreduce.py
flashinfer/comm/mnnvl.py
flashinfer/comm/trtllm_ar.py
flashinfer/comm/trtllm_mnnvl_ar.py
tests/comm/test_trtllm_allreduce_checkpoint.py
tests/comm/test_trtllm_mnnvl_allreduce.py

✅ Files skipped from review due to trivial changes (1)

docs/api/comm.rst

🚧 Files skipped from review as they are similar to previous changes (5)

tests/comm/test_trtllm_mnnvl_allreduce.py
flashinfer/comm/allreduce.py
flashinfer/comm/trtllm_mnnvl_ar.py
flashinfer/comm/mnnvl.py
flashinfer/comm/trtllm_ar.py

aleozlx · 2026-07-08T01:36:00Z

/bot run tests/comm

flashinfer-bot · 2026-07-08T01:36:36Z

GitLab MR !911 has been created, and the CI pipeline #57195830 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-07-08T06:24:55Z

[SUCCESS] Pipeline #57195830: 14/20 passed

…ore (#3727) ## Summary - Preserve the existing MNNVL CUDA virtual-address reservation while releasing and recreating its physical/imported handles across process checkpoint/restore. - Expose the lifecycle only at the owning `MoeAlltoAll` workspace as `checkpoint_prepare()` and `checkpoint_restore(comm_backend)`. - Require the restored rank-group backend to preserve the original rank and size, then remap fresh handles at the existing virtual addresses and reinitialize the MoE all-to-all workspace. - Keep handle detach/reattach private. - Treat repeated calls as successful no-ops once the workspace is already in the requested state. - Keep dispatch/combine state in each `_A2AState` and read handle attachment from the shared `MnnvlMemory.mapped` state. - Reject data operations while the MNNVL handles are detached. - Close exported, duplicated, and pidfd descriptors in the POSIX handle-exchange path. ## Scope This PR is independent of the symmetric-memory/all-reduce work in #3745. It covers the MNNVL MoE all-to-all DEP/EP path and does not add retry semantics or a second public memory-level checkpoint API. ## Usage ```python moe_alltoall.checkpoint_prepare() moe_alltoall.checkpoint_restore(comm_backend) ``` Successful transitions are idempotent. A transition that raises is terminal for that workspace; callers should abort the restored worker rather than retrying the failed operation. Both methods are collective: every rank must call them in the same order, and `comm_backend` must reproduce the original rank and world size. ## Validation Rebased onto `flashinfer-ai/main` at `2b150b39a7554fe73b5e1e8864b9f93030ebe2be`. - `ruff format --check flashinfer/comm/mnnvl.py flashinfer/comm/trtllm_moe_alltoall.py` - `ruff check flashinfer/comm/mnnvl.py flashinfer/comm/trtllm_moe_alltoall.py` - `python -m py_compile flashinfer/comm/mnnvl.py flashinfer/comm/trtllm_moe_alltoall.py` - `git diff --check` No GPU/MNNVL runtime test was run in this environment.  ## Summary by CodeRabbit * **New Features** * Added checkpoint lifecycle APIs for MoE all-to-all: `checkpoint_prepare()` and `checkpoint_restore(comm_backend)` to support CUDA checkpoint/restore with correct handle remapping, address preservation, and metadata refresh. * **Bug Fixes** * Improved MNNVL handle lifecycle management during checkpoint transitions, including safer workspace cleanup and stricter state handling. * Enforced fail-fast checks when required handles are not currently mapped. * **Documentation** * Expanded API documentation with collective call ordering, communicator matching requirements, idempotency behavior, and restart guidance after exceptions.  --------- Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>

aleozlx · 2026-07-08T19:33:38Z

hi the pre-commit check still failed

see also https://github.com/flashinfer-ai/flashinfer/blob/main/.github/pull_request_template.md#-pre-commit-checks

Wrap the little-endian Lamport communicator size in a ctypes-owned buffer before passing it to cudaMemcpy, satisfying the pre-commit mypy hook without suppressing type checking. AI-assisted: diagnosed and validated with repository pre-commit hooks. Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>

…3745) ## Summary Adds checkpoint lifecycle support for all-reduce workspaces whose pointers are captured by CUDA graphs. Physical CUDA VMM backing can be released before checkpoint and recreated at the same reserved virtual addresses after restore. The workspace API is: - `checkpoint_prepare()` collectively quiesces the group and detaches physical backing while preserving graph-visible virtual addresses. - `checkpoint_restore(comm_backend)` attaches fresh backing at those addresses, reinitializes protocol state, and collectively fences before returning. - Repeating a successfully completed prepare or restore is a no-op. - Lifecycle failures propagate without rollback or retry; the affected workspace must not be reused. - Workspaces backed by torch symmetric memory explicitly reject this lifecycle because that allocator does not provide the required stable-address remap contract. ## Implementation - Decomposes `SymmDeviceMemory` construction around `_create_and_map_handles()` and adds the inverse `_unmap_and_release_handles()` operation. Initial allocation and restore therefore use the same mapping path. - Retains CUDA virtual-address reservations while releasing unicast and multicast physical handles. - Uses a fresh communicator and handle exchanger during restore. - Supports unicast-only TRT-LLM workspaces without requiring multicast capability. - Moves the MNNVL all-reduce workspace to the same VMM-backed memory core. - Reinitializes TRT-LLM Lamport/barrier/control state and MNNVL Lamport/flag state before graph replay. - Rejects all-reduce launches while native backing is detached. ## Validation - `ruff format` leaves the changed Python files unchanged. - `ruff check` passes for the changed Python files. - `py_compile` passes for the changed Python files. - `git diff --check origin/main` passes. - The torch-symmetric-memory rejection test passes for both all-reduce workspace types. - Checkpoint and MNNVL all-reduce test collection succeeds (278 tests).  ## Summary by CodeRabbit ## Summary by CodeRabbit * **New Features** * Added `checkpoint_prepare()` / `checkpoint_restore(comm_backend)` lifecycle support for symmetric-memory–backed all-reduce fusion workspaces (TRT-LLM and MNNVL), enabling safe reattachment and protocol reinitialization. * **Bug Fixes** * Tightened validation for consistent handle attachment/detachment; fused all-reduce now rejects launches when required workspace handles are not attached. * **Documentation** * Documented required call order, rank synchronization expectations, idempotency, and unsupported backings. * **Tests** * Added distributed CUDA-graph coverage to verify behavior before/after symmetric-memory remapping.  --------- Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com> Co-authored-by: William Arnold <7565007+Aphoh@users.noreply.github.com> (cherry picked from commit 68ebdbd)

## Description vLLM can create the FlashInfer MNNVL all-reduce fusion workspace while running graph warmup under `torch.inference_mode()`. That makes `buffer_flags` an inference tensor. A later checkpoint restore can run outside inference mode and calls `_initialize_protocol()`, whose in-place `copy_` then fails with: ```text RuntimeError: Inplace update to inference tensor outside InferenceMode is not allowed. ``` This was reproduced during Dynamo Snapshot + GMS restore of GLM-5.2 with TEP8 after vLLM selected the MNNVL FlashInfer all-reduce backend. The restore failed after remapping the stable workspace VA, while reinitializing the protocol flags. Run the protocol flag reset inside a narrowly scoped `torch.inference_mode()` context. This is the mutation site that requires inference mode, so the change: - supports workspaces created either inside or outside inference mode; - does not change tensor allocation or stable-VA ownership; - does not disable inference mode or enable autograd; and - restores the caller's inference-mode state when the context exits. The regression test creates `buffer_flags` under inference mode, invokes protocol initialization from normal mode, verifies the expected flag contents, and verifies the caller remains outside inference mode afterward. ## Related Issues No associated issue. Related checkpointable MNNVL all-reduce work: #3745. ## Checklist - [x] I have added tests to cover my changes. - [x] I have run `pre-commit run --all-files` and all hooks pass. - [x] My commit includes a DCO sign-off. ## Tests Passed: ```bash PYTHONPATH="$PWD" python -m pytest \ tests/comm/test_trtllm_allreduce_checkpoint.py \ -k 'protocol_restore or lifecycle_rejects' -vv # 2 passed, 2 deselected pre-commit run --all-files # all hooks passed ``` The focused regression test also passed independently. ## Reviewer Notes The earlier workaround allocated `buffer_flags` outside inference mode. This version instead fixes the operation that has the constraint: mutating a potentially inference-created tensor during protocol restore. It avoids a separate allocation helper and avoids changing allocation semantics.  ## Summary by CodeRabbit * **Bug Fixes** * Improved protocol restoration so communication buffers are reset safely without changing the surrounding inference-mode state. * Enhanced reliability when restoring checkpointed communication state. * **Tests** * Added coverage verifying buffer flags are correctly reset during protocol restoration.  Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>

## Summary Make the existing intranode `MixedCommHandler` all-gather and reduce-scatter paths checkpointable without adding a new collective kernel. - Add one `_vm_mapped` boolean and idempotent `checkpoint_prepare()` / `checkpoint_restore(comm_backend)` methods. - Reuse the existing VMM startup structure for initial mapping and restore. - During prepare, synchronize through the retained VMM `CommBackend` and release UC/MC physical, imported, mapped, and bound resources while retaining all CUDA VA reservations and the GPU peer-pointer table captured by CUDA graphs. - During restore, validate a fresh local `CommBackend`, recreate backing resources at the exact retained addresses, reset mixed-comm protocol memory, collectively fence, and retain that backend for the next prepare or attached shutdown. - Clear the retained backend after prepare so its caller-owned process group may be destroyed while the workspace is detached. - Use the same unmap/release and permanent-address-free helpers in explicit shutdown, including shutdown from an already detached state. This PR deliberately does **not** add a standalone symmetric all-gather kernel. It makes the existing mixed-comm all-gather and reduce-scatter kernels checkpointable. ## Lifecycle contract The public API matches #3727/#3745: ```python handler.checkpoint_prepare() handler.checkpoint_restore(fresh_comm_backend) ``` The current backend must remain valid through `checkpoint_prepare()`. After prepare returns, MixedComm no longer retains it and the caller may destroy its underlying process group. The fresh backend passed to restore is retained until the next prepare or attached shutdown. MixedComm does not own or destroy externally supplied process groups. No Python launch or CUDA graph replay guard is added. The caller/checkpoint orchestrator is responsible for not executing mixed-comm operations while the workspace is detached. ## Scope Checkpoint prepare/restore currently supports intranode mixed-comm workspaces only. Multi-node checkpoint calls fail during preflight without mutating the handler. Existing multi-node/NVSHMEM execution and normal shutdown remain unchanged. `CommBackend` is control-plane rendezvous only; it does not add multi-node NVSHMEM/RDMA checkpoint support. The implementation follows the earlier checkpoint lifecycle pattern: - no additional workspace wrapper; - no attachment-state enum; - no terminal/partial-recovery state machine; - fail-fast resource transitions; - successful prepare/restore idempotence derived from `_vm_mapped`; - upstream-style unannotated dynamic CUDA resource fields. The branch is one commit directly on the PR base. Final diff: ```text flashinfer/comm/mixed_comm.py | 221 lines changed tests/comm/test_mixed_comm_checkpoint.py | 215 lines added 2 files changed, 374 insertions(+), 62 deletions(-) ``` ## Validation ### Static and local - Full pre-commit: pass, including `mypy --all-files`, Ruff lint, and Ruff format. - Python compilation and `git diff --check`: pass. - Focused checkpoint test: exactly 1 test collected; expected hardware skip on SM89 because CUDA multicast is unavailable. - Existing mixed-comm suite: exactly 6 tests collected; expected SM89 capability skips. - Independent implementation review: approved. ### B200 CUDA process checkpoint/restore Validated exact commit `1c2ead3272289c2688be8d810cc019734dcc5149` on two NVIDIA B200 GPUs in `nscale-dev` using NVIDIA's official `cuda-checkpoint` utility. Before checkpoint, both ranks passed eager execution and CUDA graph capture/replay for all four combinations: ```text ALLGATHER x FUSED_OPT_WAITS_UC ALLGATHER x FUSED_OPT_WAITS_MC REDUCESCATTER x FUSED_OPT_WAITS_UC REDUCESCATTER x FUSED_OPT_WAITS_MC ``` Each rank called `checkpoint_prepare()` twice, verified it was detached, verified the backend reference was cleared, and verified exact UC/MC/pointer-table address equality. No operation was launched or replayed while detached. An external controller drove both rank PIDs through: ```text running -> locked -> checkpointed -> locked -> running ``` While both were `checkpointed`, `nvidia-smi --query-compute-apps` returned no compute processes. After CUDA restore/unlock, each rank: 1. created a fresh Gloo process group and `TorchDistBackend`; 2. called `checkpoint_restore()` twice; 3. verified the fresh backend was retained and all addresses matched exactly; 4. replayed the original four captured graphs with changed inputs and exact AG/RS results; 5. ran a second prepare/restore lifecycle using the retained fresh backend; 6. replayed all four original graphs again with another set of changed inputs; 7. completed attached shutdown while the fresh group remained valid; 8. let the caller destroy the fresh process group afterward. Both worker return codes were zero. The successful E2E runtime was approximately 41 seconds; the clean optimized SM100a JIT build took approximately 361 seconds. The temporary test pod was deleted after evidence collection.  ## Summary by CodeRabbit * **New Features** * Added checkpoint restore for mixed communication with reusable GPU virtual-memory mappings. * Extended checkpoint restore to accept a communication backend. * Switched setup/restore/shutdown synchronization to backend-mediated coordination. * **Bug Fixes** * Improved checkpoint prepare/restore correctness, including repeated checkpoint cycles. * Preserved GPU address reservations across prepare and restore, ensuring stable teardown. * **Tests** * Added distributed CUDA tests covering checkpoint/restore, CUDA graph capture and replay, and shutdown.

flashinfer-bot added the op: comm label Jun 26, 2026

gemini-code-assist Bot reviewed Jun 26, 2026

View reviewed changes

Comment thread flashinfer/comm/allreduce.py Outdated

Comment thread flashinfer/comm/trtllm_ar.py Outdated

coderabbitai Bot reviewed Jun 26, 2026

View reviewed changes

galletas1712 force-pushed the schwinns/allreduce-graph-stable-checkpoint-adjacent branch from f9005d4 to a671244 Compare June 30, 2026 23:36

galletas1712 mentioned this pull request Jul 1, 2026

feat(mnnvl): preserve MoE all-to-all graph VAs across checkpoint restore #3727

Merged

galletas1712 force-pushed the schwinns/allreduce-graph-stable-checkpoint-adjacent branch from a671244 to d46af5d Compare July 1, 2026 01:50

galletas1712 changed the title ~~feat(comm): preserve allreduce graph VAs across checkpoint reattach~~ feat(comm): preserve all-reduce graph VAs across checkpoint restore Jul 1, 2026

galletas1712 force-pushed the schwinns/allreduce-graph-stable-checkpoint-adjacent branch from d46af5d to 7d558f2 Compare July 1, 2026 01:53

This was referenced Jul 1, 2026

feat(comm): add checkpointable symmetric all-gather galletas1712/flashinfer#3

Closed

feat(comm): add checkpointable TRT-LLM all-reduce workspaces Aphoh/flashinfer#1

Open

feat(comm): add checkpointable symmetric all-gather Aphoh/flashinfer#2

Open

galletas1712 force-pushed the schwinns/allreduce-graph-stable-checkpoint-adjacent branch 5 times, most recently from 0f2b8ec to ddf0236 Compare July 1, 2026 08:15

galletas1712 marked this pull request as ready for review July 1, 2026 08:22

galletas1712 requested review from aleozlx, bkryu, cyx-6, jimmyzho, nv-yunzheq, saltyminty, sricketts, yongwww, yyihuang and yzh119 as code owners July 1, 2026 08:22

feat(comm): preserve all-reduce graph VAs across checkpoint restore

cdd7e18

Co-authored-by: William Arnold <7565007+Aphoh@users.noreply.github.com> Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>

galletas1712 force-pushed the schwinns/allreduce-graph-stable-checkpoint-adjacent branch from ddf0236 to cdd7e18 Compare July 1, 2026 08:47

coderabbitai Bot reviewed Jul 1, 2026

View reviewed changes

Comment thread tests/comm/test_trtllm_allreduce_checkpoint.py

aleozlx approved these changes Jul 8, 2026

View reviewed changes

aleozlx added the run-ci label Jul 8, 2026

galletas1712 requested a review from qiching as a code owner July 9, 2026 17:29

galletas1712 force-pushed the schwinns/allreduce-graph-stable-checkpoint-adjacent branch from a34bbc8 to 2dc23d5 Compare July 9, 2026 17:30

aleozlx enabled auto-merge (squash) July 9, 2026 20:19

aleozlx merged commit 68ebdbd into flashinfer-ai:main Jul 9, 2026
29 of 37 checks passed

kahyunnam added the v0.6.15 label Jul 9, 2026

kangbintNV mentioned this pull request Jul 10, 2026

docs: fix 3 documentation issues from nightly doc-check #3911

Merged

jiahanc mentioned this pull request Jul 13, 2026

Revert "feat(comm): preserve all-reduce graph VAs across checkpoint restore" #3939

Closed

galletas1712 mentioned this pull request Jul 13, 2026

fix(comm): restore MNNVL flags under inference mode #3950

Merged

3 tasks

kahyunnam mentioned this pull request Jul 16, 2026

v0.6.15 release highlights draft #4006

Closed

coderabbitai Bot mentioned this pull request Jul 17, 2026

test(comm): make checkpoint worker importable #4034

Merged

3 tasks

hhzhang16 mentioned this pull request Jul 20, 2026

[TRTLLM-14715][feat] preserve MNNVL all-reduce graph VAs across restore NVIDIA/TensorRT-LLM#16635

Merged

1 task

This was referenced Jul 21, 2026

test(fused_moe): Update MoEActivationPack field accordingly #4077

Closed

feat(comm): make mixed-comm VMM workspaces checkpointable #3910

Merged

Conversation

galletas1712 commented Jun 26, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Implementation

Validation

Summary by CodeRabbit

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jun 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

galletas1712 commented Jun 26, 2026

Uh oh!

coderabbitai Bot commented Jun 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

aleozlx commented Jul 8, 2026

Uh oh!

flashinfer-bot commented Jul 8, 2026

Uh oh!

flashinfer-bot commented Jul 8, 2026

Uh oh!

aleozlx commented Jul 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

galletas1712 commented Jun 26, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jun 26, 2026 •

edited

Loading

coderabbitai Bot commented Jun 26, 2026 •

edited

Loading