fix(comm): restore MNNVL flags under inference mode

[galletas1712]

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:

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

• I have added tests to cover my changes.
• I have run pre-commit run --all-files and all hooks pass.
• My commit includes a DCO sign-off.

Tests

Passed:

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.

[coderabbitai]

Caution

Review failed

The head commit changed during the review from 7270b5362d151e79b542b8f9773c1b10eb7a556d to 243d56c.

📝 Walkthrough

Walkthrough

The protocol initializer now performs buffer_flags copying inside torch.inference_mode(). A checkpoint restoration test verifies inference mode is not left enabled and that the flags reset to the expected uint32 tensor.

Changes

Protocol restoration

Layer / File(s)	Summary
Inference-safe buffer flag reset flashinfer/comm/trtllm_mnnvl_ar.py, tests/comm/test_trtllm_allreduce_checkpoint.py	_initialize_protocol() copies buffer flags within inference mode, and the new test verifies the resulting flags and inference-mode state.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• flashinfer-ai/flashinfer#2955: Updates related MNNVL workspace buffer and flag handling.
• flashinfer-ai/flashinfer#3745: Covers related protocol re-initialization behavior during checkpoint restoration.

Suggested reviewers: samnordmann, kangbintnv, timlee0212, jiahanc, amir-19

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Title check	✅ Passed	The title clearly matches the main change: restoring MNNVL flags under inference mode.
Description check	✅ Passed	The description follows the template and includes the required sections for description, issues, checklist, tests, and reviewer notes.

✨ 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

• X
• Mastodon
• Reddit
• LinkedIn

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

[gemini-code-assist]

Code Review

This pull request wraps the buffer_flags copy operation in _initialize_protocol with torch.inference_mode() to handle cases where the workspace is created under inference mode but restored outside of it. Additionally, a unit test test_protocol_restore_resets_inference_flags has been added to verify this behavior. There are no review comments to address.

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]

🧹 Nitpick comments (1)

flashinfer/comm/trtllm_mnnvl_ar.py (1)

178-183: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider extracting the shared buffer-flags layout to avoid duplication.

The [0, 2, self.buffer_size_bytes, 0, *num_bytes_to_clear, 0] construction is now duplicated between __init__ and _initialize_protocol. Since _initialize_protocol was touched here, a small helper (e.g., _default_buffer_flags_values()) would keep both call sites in sync if the layout ever changes.

Also applies to: 249-258

🤖 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 `@flashinfer/comm/trtllm_mnnvl_ar.py` around lines 178 - 183, Extract the
duplicated buffer-flags layout into a shared helper such as
_default_buffer_flags_values(), then use it in both __init__ and
_initialize_protocol when constructing self.buffer_flags. Preserve the existing
values, dtype, and CUDA device configuration while ensuring both call sites
remain synchronized.

🤖 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 `@flashinfer/comm/trtllm_mnnvl_ar.py`:
- Around line 178-183: Extract the duplicated buffer-flags layout into a shared
helper such as _default_buffer_flags_values(), then use it in both __init__ and
_initialize_protocol when constructing self.buffer_flags. Preserve the existing
values, dtype, and CUDA device configuration while ensuring both call sites
remain synchronized.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: be85c314-97e3-4635-b1b5-6a9da4af8b44

📥 Commits

Reviewing files that changed from the base of the PR and between 78333e8 and 243d56c.

📒 Files selected for processing (2)

• flashinfer/comm/trtllm_mnnvl_ar.py
• tests/comm/test_trtllm_allreduce_checkpoint.py

[kahyunnam]

/bot run tests/gemm

[flashinfer-bot]

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

[flashinfer-bot]

[FAILED] Pipeline #58040314: 12/20 passed