[https://nvbugs/7015107][fix] Keep the Lamport all-reduce clear size in bytes

cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.cu

@@ -37,6 +37,9 @@ struct LamportComm
     {
         counter_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[0];
         flag_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[2];
+        // Bytes the previous Lamport kernel wrote to its sub-buffer. Every Lamport kernel on this workspace reads and
+        // writes it, whatever its dtype, so it stays in bytes. Readers round it up to whole 16-byte stores: the scalar
+        // MiniMax kernel writes 4-byte words.
         clear_ptr = &reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[0];
         flag_value = *flag_ptr;
         auto comm_size = reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[1];
@@ -298,7 +301,7 @@ __global__ void __launch_bounds__(1024) minimax_reduce_rms_kernel_lamport(MiniMa
     }
 #endif
     LamportComm<NRanks> comm(params.workspace, params.rank);
-    int clear_access = comm.clear_size / kElemsPerAccess<DType>;
+    int clear_access = static_cast<int>((comm.clear_size + sizeof(float4) - 1) / sizeof(float4));
     for (int idx = access_id; idx < tot_access; idx += access_stride, token_id += token_stride)
     {
         alignas(16) DType vals[kElemsPerAccess<DType>];
@@ -380,7 +383,7 @@ __global__ void __launch_bounds__(1024) minimax_reduce_rms_kernel_lamport(MiniMa
         // Clear comm buffer that previous kernel used
         reinterpret_cast<float4*>(comm.clear_buf)[idx] = clear_vec;
     }
-    comm.update(tot_tokens * NRanks * sizeof(float) / sizeof(DType));
+    comm.update(static_cast<int64_t>(tot_tokens) * NRanks * sizeof(float));
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
     if constexpr (TriggerCompletionAtEnd)
     {
@@ -682,15 +685,15 @@ __global__ void __launch_bounds__(1024) minimax_reduce_qk_rms_kernel_lamport_flo
     }
 
     // Clear comm buffer
-    int clear_access = static_cast<int>(comm.clear_size / (sizeof(float4) / sizeof(DType)));
+    int clear_access = static_cast<int>((comm.clear_size + sizeof(float4) - 1) / sizeof(float4));
     int clear_stride = group_stride * blockDim.x;
 
     for (int idx = group_id * blockDim.x + threadIdx.x; idx < clear_access; idx += clear_stride)
     {
         reinterpret_cast<float4*>(comm.clear_buf)[idx] = clear_vec;
     }
 
-    comm.update((2 * tot_groups * TokenPerBlock * sizeof(float) / sizeof(DType) * NRanks));
+    comm.update(static_cast<int64_t>(2) * tot_groups * TokenPerBlock * sizeof(float) * NRanks);
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
     if constexpr (TriggerCompletionAtEnd)
     {

cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -71,6 +71,9 @@ struct LamportComm
     {
         counter_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[0];
         flag_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[2];
+        // Bytes the previous Lamport kernel wrote to its sub-buffer. Every Lamport kernel on this workspace reads and
+        // writes it, whatever its dtype, so it stays in bytes. Readers round it up to whole 16-byte stores: the scalar
+        // MiniMax kernel writes 4-byte words.
         clear_ptr = &reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[0];
         flag_value = *flag_ptr;
         auto comm_size = reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[1];
@@ -465,7 +468,7 @@ __global__ void __launch_bounds__(1024) allreduce_fusion_kernel_oneshot_lamport(
     }
 #endif
     LamportComm<NRanks> comm(params.workspace, params.rank);
-    int clear_access = comm.clear_size / kElemsPerAccess<DType>;
+    int clear_access = static_cast<int>((comm.clear_size + sizeof(float4) - 1) / sizeof(float4));
 
     for (int idx = access_id; idx < tot_access; idx += access_stride)
     {
@@ -514,7 +517,7 @@ __global__ void __launch_bounds__(1024) allreduce_fusion_kernel_oneshot_lamport(
         fused_op(sum_val, tidx);
     }
 
-    comm.update(params.size * NRanks);
+    comm.update(static_cast<int64_t>(params.size) * NRanks * sizeof(DType));
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
     if constexpr (TriggerCompletionAtEnd)

cpp/tensorrt_llm/kernels/communicationKernels/allReduceWorkspace.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -69,7 +69,7 @@ Workspace::Workspace(int rank, int tp_size, int max_token_num, int hidden_dim,
     std::vector<int> h_flag_data{0, 0, 0};
     TLLM_CUDA_CHECK(cudaMemcpy(m_flag_d_ptr, h_flag_data.data(), 3 * sizeof(int), cudaMemcpyHostToDevice));
     workspace.push_back(m_flag_d_ptr);
-    // layout_buffer[0], clear size for next lamport kernel
+    // layout_buffer[0], clear size in bytes for next lamport kernel
     // layout_buffer[1], triple buffer offset for lamport kernel
     TLLM_CUDA_CHECK(cudaMalloc(&m_layout_d_ptr, 2 * sizeof(int64_t)));
     std::vector<int64_t> h_layout_data{0, static_cast<int64_t>(lamport_comm_size)};

cpp/tensorrt_llm/kernels/moe/communication/moeAllReduceFusionKernels.cu

@@ -32,6 +32,9 @@ struct LamportComm
     {
         counter_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[0];
         flag_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[2];
+        // Bytes the previous Lamport kernel wrote to its sub-buffer. Every Lamport kernel on this workspace reads and
+        // writes it, whatever its dtype, so it stays in bytes. Readers round it up to whole 16-byte stores: the scalar
+        // MiniMax kernel writes 4-byte words.
         clear_ptr = &reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[0];
         flag_value = *flag_ptr;
         auto comm_size = reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[1];
@@ -254,7 +257,7 @@ __global__ void moereduce_allreduce_fusion_kernel_oneshot_lamport(MoeReductionAl
     cudaGridDependencySynchronize();
     cudaTriggerProgrammaticLaunchCompletion();
     LamportComm<NRanks> comm(params.workspace, params.rank);
-    int clear_access = comm.clear_size / kElemsPerAccess;
+    int clear_access = static_cast<int>((comm.clear_size + sizeof(float4) - 1) / sizeof(float4));
 
     // * MoE related
     int threadid_in_cluster = cluster.thread_rank();
@@ -383,7 +386,7 @@ __global__ void moereduce_allreduce_fusion_kernel_oneshot_lamport(MoeReductionAl
         // * Fuse
         fused_op<ResidualOut, NormOut, QuantOut, DType>(sum_val, idx, tidx, access_id_in_token, params);
     }
-    comm.update(params.size * NRanks);
+    comm.update(static_cast<int64_t>(params.size) * NRanks * sizeof(DType));
 #endif
 }
 
@@ -530,7 +533,7 @@ __global__ void moefinalize_allreduce_fusion_kernel_oneshot_lamport(MoeFinalizeA
     cudaGridDependencySynchronize();
     cudaTriggerProgrammaticLaunchCompletion();
     LamportComm<NRanks> comm(params.workspace, params.rank);
-    int clear_access = comm.clear_size / kElemsPerAccess;
+    int clear_access = static_cast<int>((comm.clear_size + sizeof(float4) - 1) / sizeof(float4));
 
     // * MoE related
     int threadid_in_cluster = cluster.thread_rank();
@@ -681,7 +684,7 @@ __global__ void moefinalize_allreduce_fusion_kernel_oneshot_lamport(MoeFinalizeA
         // * Fuse
         fused_op<ResidualOut, NormOut, QuantOut, DType>(sum_val, idx, tidx, access_id_in_token, params);
     }
-    comm.update(params.size * NRanks);
+    comm.update(static_cast<int64_t>(params.size) * NRanks * sizeof(DType));
 
 #endif
 }

tensorrt_llm/_torch/distributed/allreduce_helper.py

@@ -138,7 +138,7 @@ def allocate_allreduce_fusion_workspace(
         # flag_buffer[1], non-lamport flag
         # flag_buffer[2], lamport flag
         flag_buffer = torch.tensor([0, 0, 0], dtype=torch.int, device="cuda")
-        # layout_buffer[0], clear size for next lamport kernel
+        # layout_buffer[0], clear size in bytes for next lamport kernel
         # layout_buffer[1], triple buffer offset for lamport kernel
         layout_buffer = torch.tensor([0, lamport_buffers_size], dtype=torch.int64, device="cuda")
 

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

@@ -25,6 +25,9 @@ l0_dgx_b200:
   - unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py::test_minimax_m3_piecewise_empty_adp_rank_preserves_caches
   # 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
+  # Lamport all-reduces of mixed dtypes and odd sizes sharing one workspace.
+  - unittest/_torch/multi_gpu/test_allreduce.py::test_allreduce_lamport_mixed_dtype
+  - unittest/_torch/multi_gpu/test_allreduce.py::test_minimax_allreduce_rms_odd_tokens
   # TP2 DSpark MLA drafter (split context projection, fused all-reduce + norm)
   # vs the single-GPU drafter.
   - unittest/_torch/multi_gpu/test_dspark_drafter_tp_fusion.py
@@ -75,7 +78,8 @@ l0_dgx_b200:
   - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[4-3-VANILLA] TIMEOUT (30)
   # Synthetic weights cover Engram head/row shards and TP/EP shared-expert collectives.
   - unittest/_torch/multi_gpu/test_engram_sharding.py::test_engram_shard_roundtrip -k "world:4"
-  # Vocab-parallel token selection over TP2/TP4 vocab shards (default all-reduce).
+  # Vocab-parallel token selection over TP2/TP4 vocab shards (default all-reduce),
+  # and the pair exchange next to large bf16 one-shot all-reduces.
   - unittest/_torch/multi_gpu/test_vocab_parallel_selection.py
   - unittest/_torch/moe/test_deepseek_v41_shared_moe.py
   - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpu_mtp_ar TIMEOUT (60)

tests/unittest/_torch/multi_gpu/test_allreduce.py

@@ -20,6 +20,7 @@
 import pytest
 import torch
 from mpi4py import MPI
+from mpi4py.futures import MPIPoolExecutor
 from utils.util import check_accuracy, skip_pre_blackwell
 
 import tensorrt_llm
@@ -28,6 +29,8 @@
                                              AllReduceParams, AllReduceStrategy,
                                              MiniMaxAllReduceRMS, MoEAllReduce,
                                              MoEAllReduceParams)
+from tensorrt_llm._torch.distributed.allreduce_helper import \
+    CustomAllReduceHelper
 from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
 from tensorrt_llm._torch.modules.rms_norm import RMSNorm
 from tensorrt_llm._torch.nccl_window_tensor_scope import \
@@ -935,6 +938,163 @@ def test_minimax_allreduce_rms_qk(mpi_pool_executor, non_contiguous_input):
         assert r is True
 
 
+# About 10 ms of GPU time: long enough for the other ranks' Lamport kernel to
+# read the late rank's part of its buffer before that rank writes it.
+LATE_RANK_SLEEP_CYCLES = 20_000_000
+
+
+@torch.inference_mode()
+def run_lamport_mixed_dtype_single_rank(tensor_parallel_size: int,
+                                        dtype_order: str) -> bool:
+    """Alternate bf16 and fp32 ONESHOT all-reduces on one mapping.
+
+    Both dtypes share the mapping's Lamport workspace, and each call clears
+    the sub-buffer of the call before it. A large call followed by small
+    calls of the other dtype must clear exactly what the large call wrote:
+    clearing more runs into the buffer a later call is using, clearing less
+    leaves stale values that a later large call reads as already arrived.
+    One rank per round starts each call late, so the other ranks read its
+    part of the buffer before it arrives and any stale value is taken.
+    """
+    rank = tensorrt_llm.mpi_rank()
+    torch.cuda.set_device(rank)
+    mapping = Mapping(world_size=tensor_parallel_size,
+                      tp_size=tensor_parallel_size,
+                      rank=rank)
+    hidden_size = 6144
+    max_bytes = CustomAllReduceHelper.max_workspace_size_auto(
+        tensor_parallel_size, support_deterministic=False)
+    allreduce = {
+        dtype: AllReduce(mapping,
+                         strategy=AllReduceStrategy.ONESHOT,
+                         dtype=dtype)
+        for dtype in (torch.bfloat16, torch.float32)
+    }
+    big_dtype, small_dtype = ((torch.bfloat16, torch.float32)
+                              if dtype_order == "bf16_then_fp32" else
+                              (torch.float32, torch.bfloat16))
+    big_rows = int(0.9 * max_bytes) // (hidden_size * big_dtype.itemsize)
+    # The last call is the large shape again in the small dtype: after an
+    # under-clear it lands on the large call's stale sub-buffer.
+    small_dtype_big_rows = int(
+        0.9 * max_bytes) // (hidden_size * small_dtype.itemsize)
+    calls = [(big_rows, big_dtype), (8, small_dtype), (8, small_dtype),
+             (small_dtype_big_rows, small_dtype)]
+
+    try:
+        # Four calls per round on a triple buffer, so each round starts on a
+        # different sub-buffer.
+        for step, (rows, dtype) in enumerate(calls * 3):
+            late_rank = step // len(calls) % tensor_parallel_size
+            # Small integers sum exactly in both dtypes, so the comparison is
+            # exact and any stale or cleared element shows up.
+            parts = [
+                torch.randint(-8,
+                              9, (rows, hidden_size),
+                              generator=torch.Generator(
+                                  device="cuda").manual_seed(100 * step + r),
+                              device="cuda").to(dtype)
+                for r in range(tensor_parallel_size)
+            ]
+            expected = torch.stack(parts).float().sum(dim=0)
+            if rank == late_rank:
+                torch.cuda._sleep(LATE_RANK_SLEEP_CYCLES)
+            output = allreduce[dtype](parts[rank])
+            torch.cuda.synchronize()
+            torch.testing.assert_close(
+                output.float(),
+                expected,
+                rtol=0,
+                atol=0,
+                msg=lambda m: f"step {step}, {rows}x{hidden_size} {dtype}: {m}")
+    except Exception:
+        traceback.print_exc()
+        raise
+    return True
+
+
+# A wrong clear size can hang the ranks inside the kernel; the thread method
+# ends the session with a failure instead of waiting on the MPI pool.
+@pytest.mark.timeout(300, method="thread")
+@pytest.mark.skipif(torch.cuda.device_count() < 2,
+                    reason="Requires at least 2 GPUs for this test")
+@pytest.mark.parametrize("dtype_order", ["bf16_then_fp32", "fp32_then_bf16"])
+@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True)
+def test_allreduce_lamport_mixed_dtype(mpi_pool_executor: MPIPoolExecutor,
+                                       dtype_order: str) -> None:
+    tensor_parallel_size = mpi_pool_executor.num_workers
+    results = mpi_pool_executor.map(
+        run_lamport_mixed_dtype_single_rank,
+        *zip(*[(tensor_parallel_size, dtype_order)] * tensor_parallel_size),
+    )
+    for r in results:
+        assert r is True
+
+
+@torch.inference_mode()
+def run_minimax_rms_odd_tokens_single_rank(tensor_parallel_size: int) -> bool:
+    """One-token MiniMax reduce-RMS calls on one mapping.
+
+    Each call writes one 4-byte variance per rank, 8 bytes at TP2, and the
+    next call must still clear it: a call three calls later reuses the
+    sub-buffer and would otherwise take the late rank's old variance as
+    already arrived. Each rank holds a constant shard whose value changes
+    with the step, so an old variance changes the normalized output.
+    """
+    rank = tensorrt_llm.mpi_rank()
+    torch.cuda.set_device(rank)
+    mapping = Mapping(world_size=tensor_parallel_size,
+                      tp_size=tensor_parallel_size,
+                      rank=rank)
+    shard = 3072
+    minimax_allreduce_rms = MiniMaxAllReduceRMS(mapping=mapping)
+    weights = torch.ones((shard, ), dtype=torch.bfloat16, device="cuda")
+    try:
+        for step in range(12):
+            late_rank = step // 4 % tensor_parallel_size
+            values = [
+                float(1 + (step + r) % 5) for r in range(tensor_parallel_size)
+            ]
+            rms = (sum(v * v for v in values) / tensor_parallel_size)**0.5
+            local = torch.full((1, shard),
+                               values[rank],
+                               dtype=torch.bfloat16,
+                               device="cuda")
+            if rank == late_rank:
+                torch.cuda._sleep(LATE_RANK_SLEEP_CYCLES)
+            output = minimax_allreduce_rms(input=local,
+                                           rms_weights=weights,
+                                           eps=1e-6)
+            torch.cuda.synchronize()
+            torch.testing.assert_close(output.float(),
+                                       torch.full((1, shard),
+                                                  values[rank] / rms,
+                                                  device="cuda"),
+                                       rtol=2e-2,
+                                       atol=0,
+                                       msg=lambda m: f"step {step}: {m}")
+    except Exception:
+        traceback.print_exc()
+        raise
+    return True
+
+
+# A missed clear can also hang the ranks; see test_allreduce_lamport_mixed_dtype.
+@pytest.mark.timeout(300, method="thread")
+@pytest.mark.skipif(torch.cuda.device_count() < 2,
+                    reason="Requires at least 2 GPUs for this test")
+@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True)
+def test_minimax_allreduce_rms_odd_tokens(
+        mpi_pool_executor: MPIPoolExecutor) -> None:
+    tensor_parallel_size = mpi_pool_executor.num_workers
+    results = mpi_pool_executor.map(
+        run_minimax_rms_odd_tokens_single_rank,
+        *zip(*[(tensor_parallel_size, )] * tensor_parallel_size),
+    )
+    for r in results:
+        assert r is True
+
+
 @torch.inference_mode()
 def run_nccl_symmetric_unregistered_input(tensor_parallel_size: int,
                                           num_tokens: int, hidden_size: int,