[None][fix] MNNVL all-reduce: keep a grown workspace's predecessor alive for captured graphs

tensorrt_llm/_torch/distributed/ops.py

@@ -99,6 +99,9 @@ class _MnnvlWorkspace(TypedDict):
     # An MPI communicator under MPI, the TP ProcessGroup under a non-MPI orchestrator (Ray).
     # None between checkpoint_prepare() and a successful checkpoint_restore().
     comm: Optional[Union[_MpiCommProtocol, "torch.distributed.ProcessGroup"]]
+    # Set once the workspace is handed out during CUDA graph capture: the captured graphs launch on its buffers and
+    # flags, so a larger workspace that replaces it must keep it alive.
+    captured: bool
 
 
 def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor:
@@ -256,12 +259,20 @@ def get_or_scale_allreduce_mnnvl_workspace(
         if not workspace["handle"].is_mapped():
             raise RuntimeError("MNNVL workspace handles are not attached")
         if workspace["buffer_size_bytes"] >= (buffer_size_bytes or 0):
-            return workspace
+            return _mark_if_capturing(workspace)
 
     workspace_lock = MNNVLAllReduce._get_allreduce_mnnvl_workspace_lock(mapping)
     with workspace_lock:
-        return _get_or_scale_allreduce_mnnvl_workspace(mapping, dtype,
-                                                       buffer_size_bytes)
+        return _mark_if_capturing(
+            _get_or_scale_allreduce_mnnvl_workspace(mapping, dtype,
+                                                    buffer_size_bytes))
+
+
+def _mark_if_capturing(workspace: _MnnvlWorkspace) -> _MnnvlWorkspace:
+    """Record that a CUDA graph being captured launches on this workspace."""
+    if not workspace["captured"] and torch.cuda.is_current_stream_capturing():
+        workspace["captured"] = True
+    return workspace
 
 
 def _get_or_scale_allreduce_mnnvl_workspace(
@@ -287,6 +298,10 @@ def _get_or_scale_allreduce_mnnvl_workspace(
 
     if mapping not in allreduce_mnnvl_workspaces or allreduce_mnnvl_workspaces[
             mapping]["buffer_size_bytes"] < (buffer_size_bytes or 0):
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError(
+                "MNNVL all-reduce workspace creation or growth during CUDA graph capture: run each shape once "
+                "outside capture first")
         # Initial buffer to be large enough to support 1024 tokens * 8192 hidden_dim
         init_buffer_size_bytes = max(1024 * 8192 * elem_size, buffer_size_bytes
                                      or 0)
@@ -347,6 +362,7 @@ def _get_or_scale_allreduce_mnnvl_workspace(
                 "buffer_flags": buffer_flags,
                 "buffer_size_bytes": buffer_size_bytes,
                 "comm": comm,
+                "captured": False,
             }
         except Exception as error:
             candidate_error = error
@@ -363,6 +379,13 @@ def _get_or_scale_allreduce_mnnvl_workspace(
         _initialize_allreduce_mnnvl_protocol(candidate_workspace)
         # Hand ownership of the communicator to the workspace.
         pending_comms.pop(mapping, None)
+        previous_workspace = allreduce_mnnvl_workspaces.get(mapping)
+        if previous_workspace is not None and previous_workspace["captured"]:
+            # CUDA graphs captured on the previous workspace keep launching on its buffers and flags. One that no
+            # capture used is released by the replacement below: every rank had synchronized its device before the
+            # protocol reset above returned, so no kernel still uses it.
+            MNNVLAllReduce.allreduce_mnnvl_retired_workspaces.setdefault(
+                mapping, []).append(previous_workspace)
         allreduce_mnnvl_workspaces[mapping] = candidate_workspace
     return allreduce_mnnvl_workspaces[mapping]
 
@@ -754,6 +777,11 @@ class MNNVLAllReduce(nn.Module):
     allreduce_mnnvl_workspaces: typing.ClassVar[dict[Mapping,
                                                      _MnnvlWorkspace]] = {}
 
+    # Workspaces a larger one replaced after a CUDA graph capture used them. Those graphs still launch on their buffers
+    # and flags, so they stay alive with the process. The checkpoint hooks cover only the current workspace.
+    allreduce_mnnvl_retired_workspaces: typing.ClassVar[dict[
+        Mapping, list[_MnnvlWorkspace]]] = {}
+
     # Communicators split for a mapping whose workspace construction has not
     # succeeded yet. Ownership moves to the workspace once it is published, so
     # an entry here is never reachable from allreduce_mnnvl_workspaces.

tests/integration/test_lists/test-db/l0_b200.yml

@@ -148,6 +148,9 @@ l0_b200:
   - unittest/_torch/modeling/test_kimi_linear_checkpoint.py
   # CPU self-test of the Kimi K3 disagg parity harness comparison logic.
   - test_kimi_k3_specdec.py::test_kimi_k3_disagg_parity_selftest
+  # ------------- distributed ---------------
+  # One process: the MNNVL workspace's capture guard raises before any MNNVL call.
+  - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_workspace_creation_refuses_graph_capture
   # ------------- modules (non-MoE) ---------------
   - unittest/_torch/modules/test_fused_add_rms_norm_quant.py
   - unittest/_torch/modules/test_fused_activation_quant.py

tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml

@@ -50,6 +50,7 @@ l0_gb200_multi_gpus:
   - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant
   - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_preserves_cuda_graph_addresses
   - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_rejects_wrong_group_membership
+  - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_workspace_growth_keeps_captured_graphs
   - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_preserves_moe_graph_addresses
   - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_engine_checkpoint_coordination
   - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_failure_is_collective_and_bounded

tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py

@@ -29,7 +29,8 @@
 from tensorrt_llm._mnnvl_utils import MnnvlMemory
 from tensorrt_llm._torch.distributed import (AllReduce, AllReduceFusionOp,
                                              AllReduceParams)
-from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce
+from tensorrt_llm._torch.distributed.ops import (
+    MNNVLAllReduce, get_or_scale_allreduce_mnnvl_workspace)
 from tensorrt_llm.functional import AllReduceStrategy
 from tensorrt_llm.mapping import Mapping
 
@@ -334,6 +335,107 @@ def mnnvl_checkpoint_worker_env(_: int):
     }
 
 
+def mnnvl_growth_graph_forward(tensor_parallel_size: int,
+                               tensor_parallel_rank: int) -> bool:
+    """A graph captured on the MNNVL workspace replays correctly after an eager call grew the workspace, a replaced
+    workspace that no capture used is not kept, and growth during capture is refused."""
+    env_names = ("TLLM_TEST_MNNVL", "TRTLLM_FORCE_MNNVL_AR")
+    previous_env = {
+        name: (name in os.environ, os.environ.get(name))
+        for name in env_names
+    }
+    tensor_parallel_rank = tensorrt_llm.mpi_rank()
+    torch.cuda.set_device(tensor_parallel_rank)
+    os.environ["TLLM_TEST_MNNVL"] = "1"
+    os.environ["TRTLLM_FORCE_MNNVL_AR"] = "1"
+    mapping = None
+    try:
+        MPI.COMM_WORLD.barrier()
+        mapping = Mapping(
+            world_size=tensor_parallel_size,
+            tp_size=tensor_parallel_size,
+            rank=tensor_parallel_rank,
+        )
+        MNNVLAllReduce.allreduce_mnnvl_workspaces.pop(mapping, None)
+        MNNVLAllReduce.allreduce_mnnvl_retired_workspaces.pop(mapping, None)
+        gc.collect()
+        MPI.COMM_WORLD.barrier()
+        expected = tensor_parallel_size * (tensor_parallel_size + 1) // 2
+        with torch.inference_mode():
+            allreduce = AllReduce(
+                mapping=mapping,
+                strategy=AllReduceStrategy.MNNVL,
+                dtype=torch.bfloat16,
+            )
+            input_ = torch.full((1, 7168),
+                                tensor_parallel_rank + 1,
+                                dtype=torch.bfloat16,
+                                device="cuda")
+            allreduce(input_)
+            torch.cuda.synchronize()
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                output = allreduce(input_)
+            graph.replay()
+            torch.cuda.synchronize()
+            torch.testing.assert_close(output,
+                                       torch.full_like(output, expected))
+
+            before = MNNVLAllReduce.allreduce_mnnvl_workspaces[mapping]
+            # Two-shot footprint 2 * 4096 * 7168 * 2 B: well past the initial 16 MiB per Lamport buffer.
+            big = torch.ones((4096, 7168), dtype=torch.bfloat16, device="cuda")
+            torch.testing.assert_close(
+                allreduce(big), torch.full_like(big, tensor_parallel_size))
+            after = MNNVLAllReduce.allreduce_mnnvl_workspaces[mapping]
+            assert after["buffer_size_bytes"] > before["buffer_size_bytes"]
+            assert MNNVLAllReduce.allreduce_mnnvl_retired_workspaces[
+                mapping] == [before]
+            assert before["handle"].is_mapped()
+            del big
+
+            for step in range(3):
+                input_.fill_(tensor_parallel_rank + 2 + step)
+                graph.replay()
+                torch.cuda.synchronize()
+                torch.testing.assert_close(
+                    output,
+                    torch.full_like(
+                        output, expected + tensor_parallel_size * (1 + step)))
+
+            # No capture used `after`: growing past it keeps only `before`, which the graph still launches on.
+            wider = torch.ones((6144, 7168),
+                               dtype=torch.bfloat16,
+                               device="cuda")
+            torch.testing.assert_close(
+                allreduce(wider), torch.full_like(wider, tensor_parallel_size))
+            assert MNNVLAllReduce.allreduce_mnnvl_workspaces[mapping][
+                "buffer_size_bytes"] > after["buffer_size_bytes"]
+            assert MNNVLAllReduce.allreduce_mnnvl_retired_workspaces[
+                mapping] == [before]
+            del wider, after
+            input_.fill_(tensor_parallel_rank + 1)
+            graph.replay()
+            torch.cuda.synchronize()
+            torch.testing.assert_close(output,
+                                       torch.full_like(output, expected))
+
+            huge = torch.ones((8192, 7168), dtype=torch.bfloat16, device="cuda")
+            with pytest.raises(RuntimeError, match="during CUDA graph capture"):
+                with torch.cuda.graph(torch.cuda.CUDAGraph()):
+                    allreduce(huge)
+        return True
+    finally:
+        if mapping is not None:
+            MNNVLAllReduce.allreduce_mnnvl_workspaces.pop(mapping, None)
+            MNNVLAllReduce.allreduce_mnnvl_retired_workspaces.pop(mapping, None)
+        gc.collect()
+        for name, (was_present, value) in previous_env.items():
+            if was_present:
+                os.environ[name] = value
+            else:
+                os.environ.pop(name, None)
+
+
 def mnnvl_checkpoint_rejects_wrong_membership(world_size: int,
                                               world_rank: int) -> bool:
     env_names = ("TLLM_TEST_MNNVL", "TRTLLM_FORCE_MNNVL_AR")
@@ -705,6 +807,37 @@ def test_mnnvl_checkpoint_rejects_wrong_group_membership(
     assert all(results)
 
 
+@pytest.mark.skipif(
+    platform.machine().lower() != "aarch64" or torch.cuda.device_count() < 2
+    or not MnnvlMemory.supports_mnnvl(),
+    reason="requires at least two GB200 GPUs with fabric-backed MNNVL",
+)
+@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True)
+def test_mnnvl_workspace_growth_keeps_captured_graphs(
+        mpi_pool_executor) -> None:
+    tensor_parallel_size = mpi_pool_executor.num_workers
+    results = mpi_pool_executor.map(
+        mnnvl_growth_graph_forward,
+        [tensor_parallel_size] * tensor_parallel_size,
+        range(tensor_parallel_size),
+    )
+    assert all(results)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+def test_mnnvl_workspace_creation_refuses_graph_capture() -> None:
+    """Creating or growing the MNNVL workspace allocates collectively, so under
+    CUDA-graph capture it raises before entering the allocation. One process,
+    a group of one: the guard runs before any MNNVL call."""
+    mapping = Mapping(world_size=1, rank=0, tp_size=1)
+    MNNVLAllReduce.allreduce_mnnvl_workspaces.pop(mapping, None)
+    graph, stream = torch.cuda.CUDAGraph(), torch.cuda.Stream()
+    with torch.cuda.graph(graph, stream=stream):
+        with pytest.raises(RuntimeError, match="during CUDA graph capture"):
+            get_or_scale_allreduce_mnnvl_workspace(mapping, torch.bfloat16)
+    assert mapping not in MNNVLAllReduce.allreduce_mnnvl_workspaces
+
+
 def _make_quant_scale(reference_norm: torch.Tensor,
                       fusion_op: AllReduceFusionOp) -> torch.Tensor:
     amax = reference_norm.abs().max().float().clamp_min(1e-6)