[None][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.
| 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.
| 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.
| 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>
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.