diff --git a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py index be6e75c40dc2..a5a7eab77850 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py @@ -489,6 +489,11 @@ def __init__( ep_size=self.ep_size, health=self.ep_group_health, ) + # Keep combine storage at a fixed offset across changing dispatch + # payloads, so a later dispatch cannot overwrite a peer's live combine. + if hidden_size is not None and dtype is not None: + self._reserve_combine_region(hidden_size, dtype) + self._workspace_lifecycle.register( self, watchdog_timeout_s=alltoall_watchdog_timeout_s, @@ -681,6 +686,49 @@ def _mnnvl_checkpoint_is_idle(self) -> bool: def _mnnvl_checkpoint_reset(self) -> None: self._dispatch_state = {"phase": "idle"} + def _reserve_combine_region(self, hidden_size: int, dtype: torch.dtype) -> int: + """Keep dispatch and combine storage disjoint across workspace reuse.""" + layout = (hidden_size, dtype.itemsize) + old_layout = self._workspace_state.get("combine_storage_layout") + if old_layout is not None: + if old_layout != layout: + raise ValueError( + "shared A2A workspace requires a stable combine shape and dtype size" + ) + return self._workspace_state["combine_storage_offset"] + + # Native combine stores a dense [EP, runtime_max_tokens, hidden] tensor. + # Match the existing allocation: reserve a second full-size region only + # for CFT-capable workspaces. Low-precision combine fits in this bound. + region_bytes = pad_up( + self.ep_size * self.max_num_tokens_per_rank * hidden_size * dtype.itemsize, 128 + ) + num_regions = 2 if self.can_use_cft_counted_writes else 1 + offset = (self.workspace_size_per_rank - num_regions * region_bytes) // 128 * 128 + aux_bytes = int(self.moe_a2a_metainfo[self.PAYLOAD_DATA_OFFSET_INDEX]) + if offset < aux_bytes: + raise ValueError("A2A workspace is too small for stable combine regions") + dispatch_end = getattr(self, "_dispatch_state", {}).get("dispatch_payload_end", aux_bytes) + if dispatch_end > offset: + raise ValueError("A2A dispatch payload overlaps the reserved combine region") + self._workspace_state["combine_storage_layout"] = layout + self._workspace_state["combine_storage_offset"] = offset + return offset + + def _check_dispatch_region(self, payloads: List[torch.Tensor], max_tokens: int) -> int: + if not 0 < max_tokens <= self.max_num_tokens_per_rank: + raise ValueError("runtime token count exceeds the configured A2A capacity") + end = int(self.moe_a2a_metainfo[self.PAYLOAD_DATA_OFFSET_INDEX]) + for payload in payloads: + end = pad_up( + end + self.ep_size * max_tokens * payload.shape[1] * payload.element_size(), 128 + ) + limit = self._workspace_state.get("combine_storage_offset", self.workspace_size_per_rank) + # Check before the native dispatch can write into any peer's memory. + if end > limit: + raise ValueError("A2A dispatch payload overlaps the reserved combine region") + return end + def dispatch( self, hidden_states: torch.Tensor, @@ -731,6 +779,7 @@ def dispatch( payloads.append(token_selected_slots) if token_final_scales is not None: payloads.append(token_final_scales) + dispatch_payload_end = self._check_dispatch_region(payloads, runtime_max_tokens_per_rank) can_use_cft_for_dispatch = _use_cft_for_dispatch_payloads( can_use_cft_for_dispatch, payloads ) @@ -777,7 +826,12 @@ def dispatch( if eplb_gathered_stats.numel() == 0: eplb_gathered_stats = None self._dispatch_state["eplb_gathered_stats"] = eplb_gathered_stats - self._dispatch_state["combine_payload_offset"] = int(combine_payload_offset) + if int(combine_payload_offset) != dispatch_payload_end: + raise RuntimeError("native A2A dispatch layout disagrees with the reserved layout") + self._dispatch_state["dispatch_payload_end"] = dispatch_payload_end + self._dispatch_state["combine_payload_offset"] = self._workspace_state.get( + "combine_storage_offset", dispatch_payload_end + ) self._dispatch_state["local_num_tokens"] = token_selected_slots.size(0) self._dispatch_state["runtime_max_tokens_per_rank"] = runtime_max_tokens_per_rank self._dispatch_state["active_rank_mask_snapshot"] = active_rank_mask_snapshot @@ -898,6 +952,9 @@ def combine( final_hidden_states, self.use_low_precision_combine, ) + combine_payload_offset = self._reserve_combine_region( + final_hidden_states.shape[-1], final_hidden_states.dtype + ) output = torch.ops.trtllm.moe_a2a_combine( final_hidden_states, int(local_num_tokens), @@ -959,6 +1016,8 @@ def get_combine_payload_tensor_in_workspace( if combine_payload_offset is None: raise RuntimeError("combine_payload_offset not found in dispatch state") + combine_payload_offset = self._reserve_combine_region(hidden_size, dtype) + self._dispatch_state["combine_payload_offset"] = combine_payload_offset result = torch.ops.trtllm.moe_a2a_get_combine_payload_tensor( self.workspace, int(self.ep_rank), diff --git a/tests/unittest/_torch/moe/multi_gpu/test_moe_a2a_workspace.py b/tests/unittest/_torch/moe/multi_gpu/test_moe_a2a_workspace.py new file mode 100644 index 000000000000..ecaa11db1a45 --- /dev/null +++ b/tests/unittest/_torch/moe/multi_gpu/test_moe_a2a_workspace.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Four-rank regression for alternating quantized/unquantized A2A payloads.""" + +import pytest +import torch +from mpi4py import MPI + +import tensorrt_llm as tllm +from tensorrt_llm._torch.moe.fused_moe.communication.nvlink_one_sided import NVLinkOneSided +from tensorrt_llm.mapping import Mapping + + +@pytest.mark.parametrize("use_cft", [False, True]) +@pytest.mark.parametrize("low_precision", [False, True]) +@pytest.mark.parametrize( + "capture,in_workspace", [(False, False), (False, True), (True, False), (True, True)] +) +def test_mixed_dispatch_layout_preserves_previous_combine( + capture, in_workspace, use_cft, low_precision +): + rank = tllm.mpi_rank() + assert tllm.mpi_world_size() == 4 + torch.cuda.set_device(rank) + hidden, capacity = 6144, 32768 + mapping = Mapping( + world_size=4, rank=rank, tp_size=4, moe_ep_size=4, gpus_per_node=4, enable_attention_dp=True + ) + first = NVLinkOneSided( + mapping, + 256, + 8, + capacity, + payload_in_workspace=in_workspace, + hidden_size=hidden, + dtype=torch.bfloat16, + can_use_cft_counted_writes=use_cft, + use_low_precision_combine=low_precision, + ) + second = NVLinkOneSided( + mapping, + 256, + 8, + capacity, + payload_in_workspace=in_workspace, + hidden_size=hidden, + dtype=torch.bfloat16, + can_use_cft_counted_writes=use_cft, + use_low_precision_combine=low_precision, + ) + assert first.workspace.data_ptr() == second.workspace.data_ptr() + assert first.use_cft_for_dispatch(32) == use_cft + assert first.use_cft_for_combine(32) == use_cft + ids = torch.tensor([0, 64, 128, 192, 1, 65, 129, 193], dtype=torch.int32, device="cuda").repeat( + capacity, 1 + ) + scales = torch.full((capacity, 8), 0.125, device="cuda") + small = torch.full((capacity, hidden // 2), 85, dtype=torch.uint8, device="cuda") + sf = torch.ones((capacity, hidden // 16), dtype=torch.uint8, device="cuda") + big = torch.full((capacity, hidden), float("nan"), dtype=torch.bfloat16, device="cuda") + expert = torch.full( + (4 * capacity, hidden), float(rank + 1), dtype=torch.bfloat16, device="cuda" + ) + lengths_list = ( + [266, 9, 19034, 9], + [266, 9, 19034, 9], + [0, 1, 32, 0], + [9, 0, capacity, 1], + [1, 9, 64, 0], + [1, 9, 64, 0], + ) + + def rounds(): + checks = [] + for i, lengths in enumerate(lengths_list): + count, maximum = lengths[rank], max(lengths) + comm = first if i % 2 == 0 else second + comm.dispatch( + small[:count] if i % 2 == 0 else big[:count], + sf[:count] if i % 2 == 0 else None, + ids[:count], + scales[:count], + lengths, + ) + if in_workspace: + payload = comm.get_combine_payload_tensor_in_workspace( + maximum, hidden, torch.bfloat16 + ) + payload.fill_(rank + 1) + else: + payload = expert[: 4 * maximum] + if rank == 2: + torch.cuda._sleep(200000) + output = comm.combine(payload) + checks.append(output.eq(10).all()) + return torch.stack(checks) + + # Initialize kernels before capture; no host synchronization inside rounds. + rounds() + torch.cuda.synchronize() + MPI.COMM_WORLD.Barrier() + records = [] + if capture: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + checks = rounds() + for _ in range(20): + graph.replay() + records.append(checks.clone()) + else: + for _ in range(20): + records.append(rounds()) + valid = torch.stack(records).cpu() + failures = (~valid).nonzero().tolist() + MPI.COMM_WORLD.Barrier() + first.destroy() + second.destroy() + assert not failures, (rank, capture, in_workspace, use_cft, low_precision, failures) diff --git a/tests/unittest/_torch/moe/test_moe_a2a_workspace.py b/tests/unittest/_torch/moe/test_moe_a2a_workspace.py new file mode 100644 index 000000000000..bddae9dd6946 --- /dev/null +++ b/tests/unittest/_torch/moe/test_moe_a2a_workspace.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Capacity checks for stable one-sided dispatch/combine workspace regions.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from tensorrt_llm._torch.moe.fused_moe.communication import nvlink_one_sided as a2a +from tensorrt_llm._torch.moe.fused_moe.communication.nvlink_one_sided import NVLinkOneSided + + +@pytest.fixture(params=[False, True], ids=["fence", "cft"]) +def comm(request): + instance = NVLinkOneSided.__new__(NVLinkOneSided) + instance._destroyed = True + instance._workspace_state = {} + instance._dispatch_state = {"phase": "idle"} + instance.ep_size = 4 + instance.max_num_tokens_per_rank = 32 + instance.can_use_cft_counted_writes = request.param + instance.workspace_size_per_rank = ( + 128 + (2 + int(request.param)) * 4 * 32 * 128 * 2 + 4 * 32 * 8 * 8 + ) + instance.PAYLOAD_DATA_OFFSET_INDEX = 0 + instance._workspace_lifecycle = SimpleNamespace(metainfo=torch.tensor([128])) + return instance + + +def test_fixed_offset_is_shared_and_independent_of_dispatch_layout(comm): + offset = comm._reserve_combine_region(128, torch.bfloat16) + for count in (1, 9, 32, 2): + for quantized in (False, True): + payloads = [ + torch.empty( + (0, 64 if quantized else 128), + dtype=torch.uint8 if quantized else torch.bfloat16, + ) + ] + if quantized: + payloads.append(torch.empty((0, 8), dtype=torch.uint8)) + payloads += [torch.empty((0, 8), dtype=torch.int32), torch.empty((0, 8))] + assert comm._check_dispatch_region(payloads, count) <= offset + assert comm._reserve_combine_region(128, torch.bfloat16) == offset + assert offset % 128 == 0 + regions = 2 if comm.can_use_cft_counted_writes else 1 + assert offset + regions * 4 * 32 * 128 * 2 <= comm.workspace_size_per_rank + + +@pytest.mark.parametrize("count", [0, -1, 33]) +def test_rejects_runtime_count_outside_capacity(comm, count): + with pytest.raises(ValueError, match="capacity"): + comm._check_dispatch_region([torch.empty((0, 128))], count) + + +def test_rejects_dispatch_overlap_before_native_write(comm): + comm._reserve_combine_region(128, torch.bfloat16) + with pytest.raises(ValueError, match="overlaps"): + comm._check_dispatch_region([torch.empty((0, 1024))], 32) + + +@pytest.mark.parametrize("hidden,dtype", [(64, torch.bfloat16), (128, torch.float32)]) +def test_rejects_shared_combine_layout_change(comm, hidden, dtype): + comm._reserve_combine_region(128, torch.bfloat16) + with pytest.raises(ValueError, match="stable combine"): + comm._reserve_combine_region(hidden, dtype) + + +def test_rejects_insufficient_combine_capacity(comm): + comm.workspace_size_per_rank = 256 + with pytest.raises(ValueError, match="too small"): + comm._reserve_combine_region(128, torch.bfloat16) + assert not comm._workspace_state + + +def test_lazy_reservation_checks_first_dispatch(comm): + comm._dispatch_state["dispatch_payload_end"] = comm.workspace_size_per_rank + with pytest.raises(ValueError, match="overlaps"): + comm._reserve_combine_region(128, torch.bfloat16) + assert not comm._workspace_state + + +def test_lazy_reservation_guards_following_dispatch(comm): + payload = torch.empty((0, 64), dtype=torch.uint8) + comm._dispatch_state["dispatch_payload_end"] = comm._check_dispatch_region([payload], 9) + offset = comm._reserve_combine_region(128, torch.bfloat16) + assert comm._check_dispatch_region([payload], 32) <= offset + with pytest.raises(ValueError, match="overlaps"): + comm._check_dispatch_region([torch.empty((0, 1024))], 32) + + +@pytest.mark.parametrize("use_cft", [False, True]) +def test_failed_reservation_does_not_publish_workspace(monkeypatch, use_cft): + def init_base(self, mapping): + self.mapping = mapping + self.ep_size = mapping.world_size + self.ep_rank = mapping.rank + + monkeypatch.setattr(a2a.Communication, "__init__", init_base) + monkeypatch.setattr(NVLinkOneSided, "_init_constants", classmethod(lambda cls: None)) + monkeypatch.setattr(NVLinkOneSided, "get_aux_data_size", staticmethod(lambda *args: 128)) + for name in ( + "PAYLOAD_DATA_OFFSET_INDEX", + "FLAG_VAL_OFFSET_INDEX", + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX", + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX", + ): + monkeypatch.setattr(NVLinkOneSided, name, 0) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACES", {}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE_REFCOUNTS", {}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", None) + for name in ("MnnvlMemory", "CftMnnvlMemory"): + monkeypatch.setattr(a2a, name, MagicMock()) + monkeypatch.setattr( + torch.ops.trtllm, + "moe_a2a_initialize", + MagicMock(return_value=torch.tensor([128])), + raising=False, + ) + lifecycle = SimpleNamespace(metainfo=torch.tensor([128]), register=MagicMock()) + monkeypatch.setattr( + a2a, + "_MnnvlAlltoAllWorkspaceLifecycle", + SimpleNamespace(get_or_create=MagicMock(return_value=lifecycle)), + ) + monkeypatch.setenv("TRTLLM_MOE_A2A_WORKSPACE_MB", "1") + monkeypatch.delenv("TRTLLM_MOE_A2A_FORCE_CFT", raising=False) + with pytest.raises(ValueError, match="too small"): + NVLinkOneSided( + SimpleNamespace(world_size=4, rank=0, has_cp_helix=lambda: False), + 256, + 8, + 1024, + hidden_size=1024, + dtype=torch.bfloat16, + can_use_cft_counted_writes=use_cft, + ) + lifecycle.register.assert_not_called() + assert not NVLinkOneSided._WORKSPACES + assert not NVLinkOneSided._WORKSPACE_REFCOUNTS + assert NVLinkOneSided._WORKSPACE is None