[TRTLLM-14715][feat] preserve native MoE A2A graph VAs across restore - #16632
Conversation
Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
📝 WalkthroughWalkthroughMNNVL memory now tracks mapped allocation records, separates virtual-address reservations from backing handles, and supports checkpoint detach/restore. MoE communication paths expose checkpoint lifecycle methods, enforce mapped workspaces, refresh metadata, and reset dispatch state. ChangesMNNVL checkpoint lifecycle
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant MoeAlltoAll
participant MnnvlMemory
participant NVLinkCommunication
participant Communicator
MoeAlltoAll->>MnnvlMemory: checkpoint_prepare()
MnnvlMemory-->>NVLinkCommunication: backing handles detached
MoeAlltoAll->>MnnvlMemory: checkpoint_restore(comm)
MnnvlMemory-->>NVLinkCommunication: fresh handles remapped
NVLinkCommunication->>Communicator: synchronize and barrier
NVLinkCommunication-->>MoeAlltoAll: communication state reset
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py (1)
703-722: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider mirroring the
_require_mapped()guard here for cross-implementation consistency.
MoeAlltoAll.get_combine_payload_tensor_in_workspace()callsself._require_mapped()before the phase check, but this NVLink one-sided variant does not. In practice thephase != "dispatched"guard already blocks reaching this while handles are detached (checkpointing requires theidlephase), so this is a defense-in-depth consistency nit rather than a live bug. Aligning both keeps the two frontends symmetric if the phase semantics ever change.♻️ Optional consistency tweak
def get_combine_payload_tensor_in_workspace( self, runtime_max_tokens_per_rank: int, hidden_size: int, dtype: torch.dtype ) -> torch.Tensor: ... + self._require_mapped() if self._dispatch_state.get("phase") != "dispatched": raise RuntimeError( "get_combine_payload_tensor_in_workspace called before a successful dispatch" )🤖 Prompt for 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. In `@tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py` around lines 703 - 722, Update get_combine_payload_tensor_in_workspace to call the existing _require_mapped() guard before checking the dispatch phase, matching MoeAlltoAll’s implementation while preserving the current phase validation and behavior.
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py`:
- Around line 703-722: Update get_combine_payload_tensor_in_workspace to call
the existing _require_mapped() guard before checking the dispatch phase,
matching MoeAlltoAll’s implementation while preserving the current phase
validation and behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: bc5812ec-a3d9-4c6c-9695-5f1fa61aa472
📒 Files selected for processing (5)
tensorrt_llm/_mnnvl_utils.pytensorrt_llm/_torch/distributed/moe_alltoall.pytensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.pytensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.pytests/unittest/_torch/test_mnnvl_memory_lifecycle.py
chienchunhung
left a comment
chienchunhung
left a comment
There was a problem hiding this comment.
Thanks for working on this. I left two lifecycle questions inline.
| return | ||
| torch.cuda.synchronize() | ||
| record.comm.barrier() | ||
| cls._unmap_and_release_handles(record) |
There was a problem hiding this comment.
Could we make this lifecycle transition fail closed?
If _unmap_and_release_handles() raises after releasing only some ranks, record.mapped remains True, so a later _require_mapped() can treat a partially unmapped workspace as usable.
The inverse can happen during restore: the allocation is marked mapped before the frontend has reinitialized its metadata and completed collective readiness.
Would it make sense to track an explicit transitioning or broken state, reject data-path access while in that state, and only mark the workspace ready after initialization and the final barrier have completed?
A focused failure-injection test for an error partway through unmap or restore would also help protect this contract.
There was a problem hiding this comment.
Added explicit allocation states in _mnnvl_utils.py in this MR hhzhang16/TensorRT-LLM#1, with data-path access allowed only in MAPPED. The shared lifecycle publishes a restored allocation only after frontend initialization, synchronization, and protocol reset, and marks failed transitions BROKEN. Included failure-injection coverage in the lifecycle tests.
There was a problem hiding this comment.
Thanks for putting this together. I checked the current PR head (db678080), and I don’t see the allocation-state changes from the follow-up MR yet—the only update after 297d098b adds _require_mapped() to the one-sided payload accessor. Could you please bring the transitioning/BROKEN state and failure-injection tests into this PR so we can review the fail-closed behavior against the code that will merge?
There was a problem hiding this comment.
Those changes are intentionally isolated in hhzhang16#1 so the shared lifecycle and failure-handling changes can be reviewed there first. Then I’ll merge them into #16632.
| The local phase and watchdog checks are defense-in-depth only; they do | ||
| not prove that every wrapper sharing this workspace is quiescent. | ||
| """ | ||
| if self._dispatch_state.get("phase") != "idle": |
There was a problem hiding this comment.
Could we move this preflight to the shared workspace owner? The phase and watchdog checks here cover only this instance, but mnnvl_mem is shared by every wrapper using the same workspace. An idle instance could therefore detach the handles while a sibling is still dispatched or has a watchdog using the workspace. The docstring correctly calls out the need for global quiescence, but I could not find a Communication, ConfigurableMoE, or executor-level path that establishes it before this call. Could we register and preflight all shared owners before the first unmap, then expose the transition through the lifecycle owner?
PS: The shared instances pattern in PR #15895 may be reusable here. A two-owner test where one remains active would help verify that no detach occurs.
There was a problem hiding this comment.
In MR hhzhang16/TensorRT-LLM#1: added _MnnvlAlltoAllWorkspaceLifecycle as the allocation-scoped owner. Both frontends register with it, and it checks every registered owner plus rank-wide readiness before the first unmap. The MR also includes two-owner coverage where one owner remains active.
There was a problem hiding this comment.
Thanks, that approach sounds aligned with what I had in mind. I checked db678080, though, and _MnnvlAlltoAllWorkspaceLifecycle, the shared-owner preflight, and the two-owner negative test are not in this PR yet. Could you please bring those changes over? As the current code stands, an idle instance can still unmap a workspace while another owner is active.
There was a problem hiding this comment.
Those changes are intentionally isolated in hhzhang16#1 so the shared lifecycle and failure-handling changes can be reviewed there first. Then I’ll merge them into #16632.
| if self._state.phase != "idle": | ||
| raise RuntimeError( | ||
| "Cannot checkpoint during an active MoE All-to-All phase") | ||
| if self._alltoall_watchdog is not None: |
There was a problem hiding this comment.
should we fail if there is a watchdog? why not just tear it down during checkpoint_prepare?
There was a problem hiding this comment.
Good call, I changed it in this MR: hhzhang16#1 so that the shared lifecycle now owns the workspace watchdog. It suspends the watchdog after readiness preflight and before unmapping, then recreates it with the same configuration after metadata refresh and collective restore readiness.
There was a problem hiding this comment.
Thanks—that shared ownership model makes sense. The current PR head still rejects checkpointing when a watchdog exists and does not suspend or recreate it. Could you please bring the lifecycle-owned watchdog transition into this PR as well, including preservation of its configuration across restore?
| cls.current_rank_stride, | ||
| cls.current_mem_offset, | ||
| ) | ||
| # all_handles_data like b'\x00\x00\x00 \x00\x00\x00\x00\x8f\xec\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\t\x00\x00\x00\x00\x00\x1d\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' # noqa: E501 |
There was a problem hiding this comment.
not sure what this comment is doing here. I know this was in the FI PR but do we need this here?
Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
Funatiq
changed the title
| cls._unmap_and_release_handles(record) | ||
| record.mem_handles = [None] * record.comm_size | ||
| record.mapped = False | ||
| record.comm.barrier() |
There was a problem hiding this comment.
Could we make the failure handling collective and bounded here?
If cuMemUnmap() fails on one rank, that rank exits before this barrier while ranks that succeeded can block here indefinitely. Restore has a similar risk if one rank fails before or during _create_and_map_handles(): peers may remain in an allgather() or the frontend barrier after some mappings are already live.
Would it make sense for the shared lifecycle owner to aggregate errors through PREPARE/COMMIT/ABORT, with a multi-rank failure-injection test after one rank has unmapped or remapped?
| del obj.ptr | ||
|
|
||
|
|
||
| def test_checkpoint_prepare_preserves_va_and_is_idempotent(memory): |
There was a problem hiding this comment.
Could we add a real multi-rank MNNVL regression test for the behavior this PR introduces?
The current tests mock all CUDA calls and use _FakeComm, so they cannot catch unmatched collectives, incorrect handle ownership, changed pointers, or unsafe CUDA graph replay. It would be helpful to exercise both one- and two-sided paths through repeated prepare/restore cycles on supported hardware, verify stable VAs/pointers and outputs, replay a captured graph, and confirm that mapped memory returns to baseline.
| return | ||
| torch.cuda.synchronize() | ||
| record.comm.barrier() | ||
| cls._unmap_and_release_handles(record) |
There was a problem hiding this comment.
Thanks for putting this together. I checked the current PR head (db678080), and I don’t see the allocation-state changes from the follow-up MR yet—the only update after 297d098b adds _require_mapped() to the one-sided payload accessor. Could you please bring the transitioning/BROKEN state and failure-injection tests into this PR so we can review the fail-closed behavior against the code that will merge?
| The local phase and watchdog checks are defense-in-depth only; they do | ||
| not prove that every wrapper sharing this workspace is quiescent. | ||
| """ | ||
| if self._dispatch_state.get("phase") != "idle": |
There was a problem hiding this comment.
Thanks, that approach sounds aligned with what I had in mind. I checked db678080, though, and _MnnvlAlltoAllWorkspaceLifecycle, the shared-owner preflight, and the two-owner negative test are not in this PR yet. Could you please bring those changes over? As the current code stands, an idle instance can still unmap a workspace while another owner is active.
| if self._state.phase != "idle": | ||
| raise RuntimeError( | ||
| "Cannot checkpoint during an active MoE All-to-All phase") | ||
| if self._alltoall_watchdog is not None: |
There was a problem hiding this comment.
Thanks—that shared ownership model makes sense. The current PR head still rejects checkpointing when a watchdog exists and does not suspend or recreate it. Could you please bring the lifecycle-owned watchdog transition into this PR as well, including preservation of its configuration across restore?
|
@chienchunhung A lot of your requests are intentionally isolated in hhzhang16#1 so the shared lifecycle and failure-handling changes can be reviewed there first. Then I’ll merge them into #16632. |
3 participants
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Description
This change ports the FlashInfer stable-VA lifecycle to TensorRT-LLM's native MoE all-to-all resources. FlashInfer implemented this lifecycle for its TRT-LLM-style MoE all-to-all workspace in flashinfer-ai/flashinfer#3727.
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.