[None][fix] Keep one-sided MoE combine workspace offset stable

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),

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)

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