From 74f1c9efdd70bc0d64759d2c42b50fd9a4cf3762 Mon Sep 17 00:00:00 2001
From: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Date: Wed, 8 Jul 2026 01:04:56 -0700
Subject: [PATCH 1/8] [None][test] stabilize AutoDeploy all-reduce strategy
 tests

Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
---
 tests/integration/test_lists/waives.txt       |  7 --
 .../custom_ops/test_ad_dist_strategies.py     | 90 +++++++++++++++++--
 2 files changed, 81 insertions(+), 16 deletions(-)

diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt
index 54ee277060a7..dfb71eeea32d 100644
--- a/tests/integration/test_lists/waives.txt
+++ b/tests/integration/test_lists/waives.txt
@@ -168,9 +168,6 @@ full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507:
 full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.5-fp8kv=True] SKIP (https://nvbugs/6390307)
 full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.9-fp8kv=True] SKIP (https://nvbugs/6390307)
 full:DGX_B200/unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" SKIP (https://nvbugs/6165866)
-full:DGX_B200/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py::test_allreduce_strategies[MIN_LATENCY] SKIP (https://nvbugs/6403920)
-full:DGX_B200/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py::test_allreduce_strategies[ONESHOT] SKIP (https://nvbugs/6403920)
-full:DGX_B200/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py::test_allreduce_strategies[TWOSHOT] SKIP (https://nvbugs/6403920)
 full:DGX_H100/unittest/_torch/attention/test_attention_backends.py::test_attention_backend[qwen2_0_5b_gqa_hd64-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6403909)
 full:GB200/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_eagle3[eagle3_one_model=False-overlap_scheduler=False] SKIP (https://nvbugs/6402500)
 full:GB200/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance-eagle3_one_model=False] SKIP (https://nvbugs/6402500)
@@ -343,10 +340,6 @@ unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py::TestWanT2VTP::test_wan_t2v_
 unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_self_attention[fp8-2-1560] SKIP (https://nvbugs/6198760)
 unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_self_attention[int8-2-1560] SKIP (https://nvbugs/6198760)
 unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py::TestWanI2VBatchGeneration::test_batch_prompt_shape SKIP (https://nvbugs/6418822)
-unittest/auto_deploy/multigpu/custom_ops SKIP (https://nvbugs/6403920)
-unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py::test_allreduce_strategies[AUTO] SKIP (https://nvbugs/6403920)
-unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py::test_allreduce_strategies[NCCL] SKIP (https://nvbugs/6403920)
-unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py::test_allreduce_strategies[SYMM_MEM] SKIP (https://nvbugs/6403920)
 unittest/bindings/test_executor_bindings.py SKIP (TRTLLM-13781: legacy TensorRT examples removed; tests to be removed in follow-up PR3)
 unittest/bindings/test_transfer_agent_bindings.py::TestNixlFunctionalTransfer::test_nixl_wait_in_progress_on_zero_timeout SKIP (https://nvbugs/6260897)
 unittest/disaggregated/test_cache_transceiver_harness.py::test_single_node_transfer SKIP (https://nvbugs/6410963)
diff --git a/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py b/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
index 3b4957d08124..a8bf19edd79b 100644
--- a/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
+++ b/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
@@ -12,11 +12,19 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import faulthandler
+import logging
+import os
 import signal
+import socket
 import subprocess
+import sys
 import tempfile
+import time
+from collections.abc import Iterator
 from contextlib import contextmanager
 from pathlib import Path
+from typing import TextIO
 
 import pytest
 import torch
@@ -44,6 +52,20 @@
 pytestmark = pytest.mark.threadleak(enabled=False)
 
 
+@contextmanager
+def _tee_tllm_errors(stream: TextIO) -> Iterator[None]:
+    """Mirror TensorRT-LLM errors to a stream that Click does not capture."""
+    handler = logging.StreamHandler(stream)
+    handler.setLevel(logging.ERROR)
+    handler.setFormatter(logging.Formatter("[allreduce-strategy-debug][tllm-error] %(message)s"))
+    tllm_logger = logging.getLogger("TRT-LLM")
+    tllm_logger.addHandler(handler)
+    try:
+        yield
+    finally:
+        tllm_logger.removeHandler(handler)
+
+
 class TimeoutError(Exception):
     """Exception raised when a test times out."""
 
@@ -204,7 +226,7 @@ def _prepare_dataset(root_dir: str, temp_dir: str, model_path_or_name: str, num_
         "SYMM_MEM",
     ],
 )
-def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy):  # noqa: F811
+def test_allreduce_strategies(capfd, llm_root, shared_dataset, allreduce_strategy):  # noqa: F811
     """Test different allreduce strategies with multi-GPU configuration making sure that there are no crashes or hangs.
 
     Configuration:
@@ -299,15 +321,65 @@ def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy):  #
             ]
         )
 
-        try:
-            with timeout(TEST_TIMEOUT_SECONDS):
-                result = runner.invoke(main, args, catch_exceptions=False)
-                assert result.exit_code == 0, f"Benchmark failed with output: {result.output}"
-        except TimeoutError as e:
-            pytest.fail(
-                f"Test timed out after {TEST_TIMEOUT_SECONDS}s for strategy {allreduce_strategy}. "
-                f"This might indicate a hang (e.g., TWOSHOT without C++ fix). Error: {e}"
+        debug_prefix = f"[allreduce-strategy-debug][{allreduce_strategy}]"
+        start_time = time.monotonic()
+        with capfd.disabled():
+            debug_stream = sys.stderr
+            print(
+                f"{debug_prefix} starting on host={socket.gethostname()} pid={os.getpid()} "
+                f"timeout={TEST_TIMEOUT_SECONDS}s",
+                flush=True,
             )
+            print(
+                f"{debug_prefix} torch={torch.__version__} cuda={torch.version.cuda} "
+                f"gpu_count={torch.cuda.device_count()} "
+                f"gpus={[torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]}",
+                flush=True,
+            )
+            print(
+                f"{debug_prefix} model_path={config['args']['model']} "
+                f"LLM_MODELS_ROOT={os.environ.get('LLM_MODELS_ROOT', '<unset>')}",
+                flush=True,
+            )
+            print(
+                f"{debug_prefix} workload=requests=10 input_tokens=128 output_tokens=128 "
+                f"tp={tp_size} max_batch_size={max_batch_size} "
+                f"max_num_tokens={max_num_tokens}",
+                flush=True,
+            )
+            print(
+                f"{debug_prefix} extra_llm_api_options:\n"
+                f"{Path(extra_llm_api_options_path).read_text()}",
+                flush=True,
+            )
+
+            # Preserve useful parent-process stacks in CI even if a native CUDA/MPI wait
+            # prevents the Python-level timeout from unwinding the benchmark.
+            faulthandler.dump_traceback_later(60, repeat=True)
+            try:
+                with _tee_tllm_errors(debug_stream):
+                    with timeout(TEST_TIMEOUT_SECONDS):
+                        result = runner.invoke(main, args, catch_exceptions=False)
+                        print(
+                            f"{debug_prefix} benchmark returned exit_code={result.exit_code} "
+                            f"elapsed={time.monotonic() - start_time:.1f}s",
+                            flush=True,
+                        )
+                        assert result.exit_code == 0, (
+                            f"Benchmark failed with output: {result.output}"
+                        )
+            except TimeoutError as e:
+                pytest.fail(
+                    f"Test timed out after {TEST_TIMEOUT_SECONDS}s for strategy "
+                    f"{allreduce_strategy}. This might indicate a hang "
+                    f"(e.g., TWOSHOT without C++ fix). Error: {e}"
+                )
+            finally:
+                faulthandler.cancel_dump_traceback_later()
+                print(
+                    f"{debug_prefix} finished elapsed={time.monotonic() - start_time:.1f}s",
+                    flush=True,
+                )
 
 
 @pytest.mark.parametrize(

From a616b730f08434a3811a3a33050013419dc85faa Mon Sep 17 00:00:00 2001
From: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Date: Wed, 8 Jul 2026 07:29:02 -0700
Subject: [PATCH 2/8] [None][test] expose TWOSHOT native dispatch diagnostics

Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
---
 .../custom_ops/test_ad_dist_strategies.py     | 29 ++++++++++++-------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py b/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
index a8bf19edd79b..64684c4f6e5d 100644
--- a/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
+++ b/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
@@ -358,16 +358,25 @@ def test_allreduce_strategies(capfd, llm_root, shared_dataset, allreduce_strateg
             faulthandler.dump_traceback_later(60, repeat=True)
             try:
                 with _tee_tllm_errors(debug_stream):
-                    with timeout(TEST_TIMEOUT_SECONDS):
-                        result = runner.invoke(main, args, catch_exceptions=False)
-                        print(
-                            f"{debug_prefix} benchmark returned exit_code={result.exit_code} "
-                            f"elapsed={time.monotonic() - start_time:.1f}s",
-                            flush=True,
-                        )
-                        assert result.exit_code == 0, (
-                            f"Benchmark failed with output: {result.output}"
-                        )
+                    old_log_level = os.environ.get("TLLM_LOG_LEVEL")
+                    if allreduce_strategy == "TWOSHOT":
+                        os.environ["TLLM_LOG_LEVEL"] = "DEBUG"
+                    try:
+                        with timeout(TEST_TIMEOUT_SECONDS):
+                            result = runner.invoke(main, args, catch_exceptions=False)
+                            print(
+                                f"{debug_prefix} benchmark returned exit_code={result.exit_code} "
+                                f"elapsed={time.monotonic() - start_time:.1f}s",
+                                flush=True,
+                            )
+                            assert result.exit_code == 0, (
+                                f"Benchmark failed with output: {result.output}"
+                            )
+                    finally:
+                        if old_log_level is None:
+                            os.environ.pop("TLLM_LOG_LEVEL", None)
+                        else:
+                            os.environ["TLLM_LOG_LEVEL"] = old_log_level
             except TimeoutError as e:
                 pytest.fail(
                     f"Test timed out after {TEST_TIMEOUT_SECONDS}s for strategy "

From bbcc8cfc7792d3a68d365b3c0427e1895380f141 Mon Sep 17 00:00:00 2001
From: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Date: Wed, 8 Jul 2026 07:36:58 -0700
Subject: [PATCH 3/8] [None][fix] handle partial-warp fused all-reduce RMSNorm

Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
---
 .../allReduceFusionKernels.cu                 | 28 +++++++++++++++++--
 .../kernels/allReduce/allReduceKernelTest.cu  | 12 +++++++-
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
index 5a3edda04a70..2574e1b63e71 100644
--- a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
+++ b/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.
@@ -24,6 +24,23 @@ TRTLLM_NAMESPACE_BEGIN
 
 namespace kernels::ar_fusion
 {
+__device__ __forceinline__ float warpReduceSumActive(float val)
+{
+    auto const mask = __activemask();
+    auto const lane = threadIdx.x & (warpSize - 1);
+    auto const activeLanes = __popc(mask);
+#pragma unroll
+    for (int offset = warpSize / 2; offset > 0; offset >>= 1)
+    {
+        auto const other = __shfl_down_sync(mask, val, offset);
+        if (lane + offset < activeLanes)
+        {
+            val += other;
+        }
+    }
+    return val;
+}
+
 template <int NRanks>
 struct SyncComm
 {
@@ -296,7 +313,14 @@ protected:
             float v = static_cast<float>(reinterpret_cast<DType const*>(&residual)[i]);
             acc += v * v;
         }
-        tensorrt_llm::common::blockReduceSumV2<float, 1>(&acc);
+        if (blockDim.x < warpSize)
+        {
+            acc = warpReduceSumActive(acc);
+        }
+        else
+        {
+            tensorrt_llm::common::blockReduceSumV2<float, 1>(&acc);
+        }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
         cg::cluster_group cluster = cg::this_cluster();
         if (cluster.num_blocks() > 1)
diff --git a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu
index b3d120e7015c..f2e01e6c5048 100644
--- a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu
+++ b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, 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.
@@ -620,6 +620,16 @@ TEST(Kernel, AllReduce)
         }
     }
     EXPECT_TRUE(pass);
+
+    // Exercise the partial-warp RMSNorm reduction used by small AutoDeploy test models.
+    // At hidden_size=64, a 16-bit dtype processes eight 128-bit accesses per token, so the
+    // all-reduce kernel launches fewer than one warp of threads.
+    Workspace small_hidden_workspace(world_size, rank, 8192, 64);
+    pass = pass
+        && test(small_hidden_workspace, 8192, 64, false, true, 3, 3, AllReduceStrategyType::TWOSHOT,
+            AllReduceStrategyConfig(0), AllReduceFusionOp::RESIDUAL_RMS_NORM);
+    EXPECT_TRUE(pass);
+
     ops[0] = AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM;
     for (auto config : configs)
     {

From 700dd9262f0e7d326d6cd9c2b28b6f4ef3062886 Mon Sep 17 00:00:00 2001
From: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Date: Wed, 8 Jul 2026 21:44:01 -0700
Subject: [PATCH 4/8] [None][chore] remove all-reduce CI diagnostics

The fused all-reduce residual/RMSNorm kernel derives its block size from the number of 128-bit accesses. For the AutoDeploy test model, hidden_size=64 with a 16-bit dtype produces an eight-thread block. blockReduceSumV2 assumes a complete warp, so its shuffle named inactive lanes and triggered undefined CUDA behavior. On B200 this caused the TWOSHOT CUDA-graph warmup to hang.

The production fix uses an active-mask warp reduction for sub-warp blocks while retaining the existing reduction for full-warp and multi-warp launches. The native regression test covers the exact 8192 x 64 TWOSHOT fused-residual-RMSNorm geometry, and the complete B200 custom-ops directory now passes.

Remove the temporary fault-handler, logger tee, environment logging, and forced TWOSHOT debug level used to diagnose the CI hang.

Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
---
 .../custom_ops/test_ad_dist_strategies.py     | 99 ++-----------------
 1 file changed, 9 insertions(+), 90 deletions(-)

diff --git a/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py b/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
index 64684c4f6e5d..3b4957d08124 100644
--- a/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
+++ b/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
@@ -12,19 +12,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import faulthandler
-import logging
-import os
 import signal
-import socket
 import subprocess
-import sys
 import tempfile
-import time
-from collections.abc import Iterator
 from contextlib import contextmanager
 from pathlib import Path
-from typing import TextIO
 
 import pytest
 import torch
@@ -52,20 +44,6 @@
 pytestmark = pytest.mark.threadleak(enabled=False)
 
 
-@contextmanager
-def _tee_tllm_errors(stream: TextIO) -> Iterator[None]:
-    """Mirror TensorRT-LLM errors to a stream that Click does not capture."""
-    handler = logging.StreamHandler(stream)
-    handler.setLevel(logging.ERROR)
-    handler.setFormatter(logging.Formatter("[allreduce-strategy-debug][tllm-error] %(message)s"))
-    tllm_logger = logging.getLogger("TRT-LLM")
-    tllm_logger.addHandler(handler)
-    try:
-        yield
-    finally:
-        tllm_logger.removeHandler(handler)
-
-
 class TimeoutError(Exception):
     """Exception raised when a test times out."""
 
@@ -226,7 +204,7 @@ def _prepare_dataset(root_dir: str, temp_dir: str, model_path_or_name: str, num_
         "SYMM_MEM",
     ],
 )
-def test_allreduce_strategies(capfd, llm_root, shared_dataset, allreduce_strategy):  # noqa: F811
+def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy):  # noqa: F811
     """Test different allreduce strategies with multi-GPU configuration making sure that there are no crashes or hangs.
 
     Configuration:
@@ -321,74 +299,15 @@ def test_allreduce_strategies(capfd, llm_root, shared_dataset, allreduce_strateg
             ]
         )
 
-        debug_prefix = f"[allreduce-strategy-debug][{allreduce_strategy}]"
-        start_time = time.monotonic()
-        with capfd.disabled():
-            debug_stream = sys.stderr
-            print(
-                f"{debug_prefix} starting on host={socket.gethostname()} pid={os.getpid()} "
-                f"timeout={TEST_TIMEOUT_SECONDS}s",
-                flush=True,
+        try:
+            with timeout(TEST_TIMEOUT_SECONDS):
+                result = runner.invoke(main, args, catch_exceptions=False)
+                assert result.exit_code == 0, f"Benchmark failed with output: {result.output}"
+        except TimeoutError as e:
+            pytest.fail(
+                f"Test timed out after {TEST_TIMEOUT_SECONDS}s for strategy {allreduce_strategy}. "
+                f"This might indicate a hang (e.g., TWOSHOT without C++ fix). Error: {e}"
             )
-            print(
-                f"{debug_prefix} torch={torch.__version__} cuda={torch.version.cuda} "
-                f"gpu_count={torch.cuda.device_count()} "
-                f"gpus={[torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]}",
-                flush=True,
-            )
-            print(
-                f"{debug_prefix} model_path={config['args']['model']} "
-                f"LLM_MODELS_ROOT={os.environ.get('LLM_MODELS_ROOT', '<unset>')}",
-                flush=True,
-            )
-            print(
-                f"{debug_prefix} workload=requests=10 input_tokens=128 output_tokens=128 "
-                f"tp={tp_size} max_batch_size={max_batch_size} "
-                f"max_num_tokens={max_num_tokens}",
-                flush=True,
-            )
-            print(
-                f"{debug_prefix} extra_llm_api_options:\n"
-                f"{Path(extra_llm_api_options_path).read_text()}",
-                flush=True,
-            )
-
-            # Preserve useful parent-process stacks in CI even if a native CUDA/MPI wait
-            # prevents the Python-level timeout from unwinding the benchmark.
-            faulthandler.dump_traceback_later(60, repeat=True)
-            try:
-                with _tee_tllm_errors(debug_stream):
-                    old_log_level = os.environ.get("TLLM_LOG_LEVEL")
-                    if allreduce_strategy == "TWOSHOT":
-                        os.environ["TLLM_LOG_LEVEL"] = "DEBUG"
-                    try:
-                        with timeout(TEST_TIMEOUT_SECONDS):
-                            result = runner.invoke(main, args, catch_exceptions=False)
-                            print(
-                                f"{debug_prefix} benchmark returned exit_code={result.exit_code} "
-                                f"elapsed={time.monotonic() - start_time:.1f}s",
-                                flush=True,
-                            )
-                            assert result.exit_code == 0, (
-                                f"Benchmark failed with output: {result.output}"
-                            )
-                    finally:
-                        if old_log_level is None:
-                            os.environ.pop("TLLM_LOG_LEVEL", None)
-                        else:
-                            os.environ["TLLM_LOG_LEVEL"] = old_log_level
-            except TimeoutError as e:
-                pytest.fail(
-                    f"Test timed out after {TEST_TIMEOUT_SECONDS}s for strategy "
-                    f"{allreduce_strategy}. This might indicate a hang "
-                    f"(e.g., TWOSHOT without C++ fix). Error: {e}"
-                )
-            finally:
-                faulthandler.cancel_dump_traceback_later()
-                print(
-                    f"{debug_prefix} finished elapsed={time.monotonic() - start_time:.1f}s",
-                    flush=True,
-                )
 
 
 @pytest.mark.parametrize(

From 8c530c72b2fa04c6f1641bcdb52b035af3a77aaf Mon Sep 17 00:00:00 2001
From: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Date: Sat, 11 Jul 2026 23:18:41 -0700
Subject: [PATCH 5/8] fixed

Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
---
 .../allReduceFusionKernels.cu                 | 50 +++++++++++++++----
 .../kernels/allReduce/allReduceFusionTest.cu  | 14 +++---
 2 files changed, 47 insertions(+), 17 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
index 2574e1b63e71..397bc0032b3f 100644
--- a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
+++ b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
@@ -24,23 +24,55 @@ TRTLLM_NAMESPACE_BEGIN
 
 namespace kernels::ar_fusion
 {
-__device__ __forceinline__ float warpReduceSumActive(float val)
+__device__ __forceinline__ uint32_t getLaneMask(int activeLanes)
 {
-    auto const mask = __activemask();
-    auto const lane = threadIdx.x & (warpSize - 1);
-    auto const activeLanes = __popc(mask);
+    return activeLanes == warpSize ? 0xffffffffU : ((1U << activeLanes) - 1U);
+}
+
+__device__ __forceinline__ float warpReduceSum(float val, uint32_t mask, int activeLanes)
+{
+    int const lane = threadIdx.x & (warpSize - 1);
 #pragma unroll
     for (int offset = warpSize / 2; offset > 0; offset >>= 1)
     {
-        auto const other = __shfl_down_sync(mask, val, offset);
-        if (lane + offset < activeLanes)
+        if (offset < activeLanes)
         {
-            val += other;
+            float const other = __shfl_xor_sync(mask, val, offset, warpSize);
+            if ((lane ^ offset) < activeLanes)
+            {
+                val += other;
+            }
         }
     }
     return val;
 }
 
+__device__ __forceinline__ void blockReduceSum(float* val)
+{
+    __shared__ float shared[32];
+    int const lane = threadIdx.x & (warpSize - 1);
+    int const warpId = threadIdx.x / warpSize;
+    int const warpStart = warpId * warpSize;
+    int const remainingLanes = static_cast<int>(blockDim.x) - warpStart;
+    int const activeLanes = remainingLanes < warpSize ? remainingLanes : warpSize;
+    uint32_t const mask = getLaneMask(activeLanes);
+
+    *val = warpReduceSum(*val, mask, activeLanes);
+    if (lane == 0)
+    {
+        shared[warpId] = *val;
+    }
+
+    __syncthreads();
+
+    int const warpCount = (blockDim.x + warpSize - 1) / warpSize;
+    if (warpCount > 1 && warpId == 0)
+    {
+        *val = lane < warpCount ? shared[lane] : 0.f;
+        *val = warpReduceSum(*val, 0xffffffffU, warpSize);
+    }
+}
+
 template <int NRanks>
 struct SyncComm
 {
@@ -313,9 +345,9 @@ protected:
             float v = static_cast<float>(reinterpret_cast<DType const*>(&residual)[i]);
             acc += v * v;
         }
-        if (blockDim.x < warpSize)
+        if (blockDim.x % warpSize != 0)
         {
-            acc = warpReduceSumActive(acc);
+            blockReduceSum(&acc);
         }
         else
         {
diff --git a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.cu b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.cu
index 4b9c7af29a46..661e9cee1b26 100644
--- a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.cu
+++ b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.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.
@@ -429,10 +429,8 @@ public:
         }
         else if constexpr (ar_fusion::GetQuantType<Pattern> == ar_fusion::QuantType::kFP8)
         {
-            // We need norm out to verify the accuracy of quantization.
-            static_assert(Pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP8Quant);
             CudaBuffer ref_fp8_output(message_size * sizeof(__nv_fp8_e4m3));
-            quantize_to_fp8(m_norm_out.device_data<DType>(), ref_fp8_output.device_data<__nv_fp8_e4m3>(), message_size,
+            quantize_to_fp8(ref_output.device_data<DType>(), ref_fp8_output.device_data<__nv_fp8_e4m3>(), message_size,
                 m_scale_factor.device_data<float>(), m_stream->get());
             TLLM_CHECK(compare<__nv_fp8_e4m3>(
                 m_rank, m_quant_out.host_data(), ref_fp8_output.host_data(), message_size, "fp8 quant out", 0));
@@ -600,14 +598,14 @@ TEST(Kernel_AllReduceFusion, AllReduceFusionAccuracyDifferentHiddenDim)
     ASSERT_EQ(world_size % 2, 0) << "Requires even world size (got " << world_size << ")";
 
     int const arch = tensorrt_llm::common::getSMVersion();
+    TEST_AR_FUSION(half, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP8Quant);
+    TEST_AR_FUSION(half, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP8Quant);
+    TEST_AR_FUSION(__nv_bfloat16, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP8Quant);
+    TEST_AR_FUSION(__nv_bfloat16, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP8Quant);
     if (arch >= 100)
     {
         TEST_AR_FUSION(half, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP4Quant);
     }
-    else
-    {
-        TEST_AR_FUSION(half, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP8Quant);
-    }
 #undef TEST_AR_FUSION
 }
 

From d39989e96f4f2b9e10198892f0798a804efc6a28 Mon Sep 17 00:00:00 2001
From: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Date: Sun, 12 Jul 2026 04:44:02 -0700
Subject: [PATCH 6/8] moved test to smoke

Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
---
 .../custom_ops/test_ad_dist_strategies.py     | 276 ----------------
 .../smoke/test_ad_allreduce_strategies.py     | 297 ++++++++++++++++++
 2 files changed, 297 insertions(+), 276 deletions(-)
 create mode 100644 tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py

diff --git a/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py b/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
index 3b4957d08124..b9ae83d8f020 100644
--- a/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
+++ b/tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
@@ -12,18 +12,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import signal
-import subprocess
-import tempfile
-from contextlib import contextmanager
-from pathlib import Path
-
 import pytest
 import torch
 import torch.nn as nn
-import yaml
-from _model_test_utils import get_small_model_config
-from click.testing import CliRunner
 from utils.cpp_paths import llm_root  # noqa: F401
 
 from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm
@@ -37,279 +28,12 @@
 from tensorrt_llm._torch.auto_deploy.transform.library.sharding_ir import SplitDimension
 from tensorrt_llm._torch.auto_deploy.utils._graph import recompile
 from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op
-from tensorrt_llm.commands.bench import main
 from tensorrt_llm.functional import AllReduceStrategy
 
 # needed since LLM API uses MPI executor pool internally for TP>1, which leaks a thread on shutdown
 pytestmark = pytest.mark.threadleak(enabled=False)
 
 
-class TimeoutError(Exception):
-    """Exception raised when a test times out."""
-
-    pass
-
-
-@contextmanager
-def timeout(seconds):
-    """Context manager that raises TimeoutError if code block exceeds time limit.
-
-    Args:
-        seconds: Maximum time in seconds to allow the code block to run
-
-    Raises:
-        TimeoutError: If the code block execution exceeds the time limit
-    """
-
-    def timeout_handler(signum, frame):
-        raise TimeoutError(f"Test execution exceeded {seconds} seconds timeout")
-
-    # Set the signal handler and alarm
-    old_handler = signal.signal(signal.SIGALRM, timeout_handler)
-    signal.alarm(seconds)
-    try:
-        yield
-    finally:
-        # Restore the old signal handler and cancel the alarm
-        signal.alarm(0)
-        signal.signal(signal.SIGALRM, old_handler)
-
-
-@pytest.fixture(scope="module", autouse=True)
-def prewarm_flashinfer_jit():
-    """Pre-warm FlashInfer JIT kernels before multi-GPU tests.
-
-    This prevents a race condition where multiple MPI ranks try to JIT-compile
-    FlashInfer kernels simultaneously to the same cache directory, causing
-    Ninja build failures like: "ninja: error: opening build log: No such file or directory"
-
-    By triggering the compilation in the main process first, the kernels are
-    cached and available for all worker ranks.
-    """
-    try:
-        import flashinfer
-        import flashinfer.page
-        import flashinfer.sampling
-
-        if torch.cuda.is_available():
-            # Prevent concurrent JIT warmup across multiple pytest processes (e.g., xdist).
-            try:
-                import fcntl  # Linux-only
-            except ImportError:
-                fcntl = None
-
-            lock_f = None
-            if fcntl is not None:
-                import pathlib
-                import tempfile
-
-                lock_path = pathlib.Path(tempfile.gettempdir()) / "flashinfer_jit_prewarm.lock"
-                lock_f = open(lock_path, "w")
-                fcntl.flock(lock_f.fileno(), fcntl.LOCK_EX)
-            # Create dummy tensors to trigger kernel JIT compilation
-            with torch.no_grad():
-                device = torch.device("cuda:0")
-
-                # Trigger page kernel compilation
-                try:
-                    # Force module loading (this triggers JIT compilation)
-                    _ = flashinfer.page.gen_page_module()
-                except Exception as exc:  # noqa: BLE001
-                    import warnings
-
-                    warnings.warn(f"FlashInfer page-kernel prewarm failed: {exc!r}", RuntimeWarning)
-
-                # Trigger sampling kernel compilation
-                try:
-                    dummy_probs = torch.softmax(torch.randn(1, 100, device=device), dim=-1)
-                    _ = flashinfer.sampling.sampling_from_probs(dummy_probs, deterministic=True)
-                except Exception as exc:  # noqa: BLE001
-                    import warnings
-
-                    warnings.warn(
-                        f"FlashInfer sampling-kernel prewarm failed: {exc!r}", RuntimeWarning
-                    )
-
-                torch.cuda.empty_cache()
-            if lock_f is not None:
-                lock_f.close()
-
-    except ImportError:
-        pass  # FlashInfer not available
-
-    yield
-
-
-@pytest.fixture(scope="module")
-def shared_dataset(llm_root):  # noqa: F811
-    """Prepare dataset once for all tests in this module."""
-    model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-    config = get_small_model_config(model_name)
-    with tempfile.TemporaryDirectory() as temp_dir:
-        dataset_path = _prepare_dataset(
-            llm_root, temp_dir, config["args"]["model"], num_requests=10
-        )
-        # Read dataset content to return it (temp_dir will be deleted)
-        with open(dataset_path, "r") as f:
-            dataset_content = f.read()
-        yield dataset_content
-
-
-def _prepare_dataset(root_dir: str, temp_dir: str, model_path_or_name: str, num_requests: int = 10):
-    """Prepare a synthetic dataset for benchmarking."""
-    _DATASET_NAME = "synthetic_128_128.txt"
-    dataset_path = Path(temp_dir, _DATASET_NAME)
-    dataset_tool = Path(root_dir, "benchmarks", "cpp", "prepare_dataset.py")
-    script_dir = Path(root_dir, "benchmarks", "cpp")
-
-    # Generate a small dataset to run a test - matching workload configuration
-    command = [
-        "python3",
-        f"{dataset_tool}",
-        "--stdout",
-        "--tokenizer",
-        model_path_or_name,
-        "token-norm-dist",
-        "--input-mean",
-        "128",
-        "--output-mean",
-        "128",
-        "--input-stdev",
-        "0",
-        "--output-stdev",
-        "0",
-        "--num-requests",
-        str(num_requests),
-    ]
-    print(f"Running command: {' '.join(command)}")
-    result = subprocess.run(
-        command, cwd=str(script_dir), capture_output=True, text=True, timeout=300
-    )
-    if result.returncode != 0:
-        raise RuntimeError(f"Failed to prepare dataset: {result.stderr}")
-    # Grab the stdout and write it to a dataset file for passing to suite.
-    with open(dataset_path, "w") as dataset:
-        dataset.write(result.stdout)
-    return dataset_path
-
-
-@pytest.mark.parametrize(
-    "allreduce_strategy",
-    [
-        "AUTO",
-        "ONESHOT",
-        "TWOSHOT",
-        "MIN_LATENCY",
-        "NCCL",
-        "SYMM_MEM",
-    ],
-)
-def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy):  # noqa: F811
-    """Test different allreduce strategies with multi-GPU configuration making sure that there are no crashes or hangs.
-
-    Configuration:
-        The allreduce_strategy is set in the transforms config:
-        ```yaml
-        transforms:
-          detect_sharding:
-            allreduce_strategy: "ONESHOT"  # or AUTO, NCCL, TWOSHOT, etc.
-        ```
-
-    Test configuration:
-        - Model: Llama-3.1-8B with TP=2
-        - Dataset: 10 synthetic requests (128 input, 128 output tokens)
-        - Timeout: 300 seconds to catch hangs
-        - Skipped if fewer than 2 GPUs available
-
-    Args:
-        llm_root: Root directory fixture
-        shared_dataset: Shared dataset fixture (prepared once for all test runs)
-        allreduce_strategy: Strategy to test (AUTO, ONESHOT, TWOSHOT, MIN_LATENCY, NCCL)
-    """
-    # Fixed timeout for all strategies (5 minutes should be enough)
-    TEST_TIMEOUT_SECONDS = 300
-
-    model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-    config = get_small_model_config(model_name)
-    tp_size = 2
-    max_batch_size = 256
-    max_num_tokens = 8192
-
-    if not torch.cuda.is_available() or torch.cuda.device_count() < tp_size:
-        pytest.skip(f"Allreduce strategy test requires at least {tp_size} GPUs, skipping")
-
-    with tempfile.TemporaryDirectory() as temp_dir:
-        # Write shared dataset to temp location
-        dataset_path = Path(temp_dir, "synthetic_128_128.txt")
-        with open(dataset_path, "w") as f:
-            f.write(shared_dataset)
-
-        # Create configuration with specified allreduce strategy in transforms
-        extra_llm_api_options_path = f"{temp_dir}/extra_llm_api_options.yaml"
-        with open(extra_llm_api_options_path, "w") as f:
-            yaml.dump(
-                {
-                    **config["args"],
-                    "max_batch_size": max_batch_size,
-                    "max_num_tokens": max_num_tokens,
-                    "max_seq_len": 256,
-                    "transforms": {
-                        "detect_sharding": {
-                            "stage": "sharding",
-                            "allreduce_strategy": allreduce_strategy,
-                        },
-                        "compile_model": {
-                            "stage": "compile",
-                            "backend": "torch-cudagraph",
-                            "cuda_graph_batch_sizes": [1, 2, 4, 8, 16, 32, 64, 128, 256],
-                        },
-                    },
-                },
-                f,
-            )
-
-        # Run benchmark with specified allreduce strategy with timeout protection
-        runner = CliRunner()
-        args = [
-            "--model",
-            model_name,
-        ]
-
-        # Only pass --model_path if it's a local filesystem path
-        # Note: --model_path must come BEFORE the subcommand (throughput)
-        if str(config["args"]["model"]).startswith("/"):
-            args.extend(["--model_path", str(config["args"]["model"])])
-
-        # Add the subcommand and its options
-        args.extend(
-            [
-                "throughput",
-                "--backend",
-                "_autodeploy",
-                "--dataset",
-                str(dataset_path),
-                "--extra_llm_api_options",
-                extra_llm_api_options_path,
-                "--tp",
-                str(tp_size),
-                "--max_batch_size",
-                str(max_batch_size),
-                "--max_num_tokens",
-                str(max_num_tokens),
-            ]
-        )
-
-        try:
-            with timeout(TEST_TIMEOUT_SECONDS):
-                result = runner.invoke(main, args, catch_exceptions=False)
-                assert result.exit_code == 0, f"Benchmark failed with output: {result.output}"
-        except TimeoutError as e:
-            pytest.fail(
-                f"Test timed out after {TEST_TIMEOUT_SECONDS}s for strategy {allreduce_strategy}. "
-                f"This might indicate a hang (e.g., TWOSHOT without C++ fix). Error: {e}"
-            )
-
-
 @pytest.mark.parametrize(
     "strategy",
     [
diff --git a/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py
new file mode 100644
index 000000000000..937bdd4c3ff9
--- /dev/null
+++ b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py
@@ -0,0 +1,297 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import signal
+import subprocess
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+
+import pytest
+import torch
+import yaml
+from _model_test_utils import get_small_model_config
+from click.testing import CliRunner
+from utils.cpp_paths import llm_root  # noqa: F401
+
+from tensorrt_llm.commands.bench import main
+
+# needed since LLM API uses MPI executor pool internally for TP>1, which leaks a thread on shutdown
+pytestmark = pytest.mark.threadleak(enabled=False)
+
+
+class TimeoutError(Exception):
+    """Exception raised when a test times out."""
+
+    pass
+
+
+@contextmanager
+def timeout(seconds):
+    """Context manager that raises TimeoutError if code block exceeds time limit.
+
+    Args:
+        seconds: Maximum time in seconds to allow the code block to run
+
+    Raises:
+        TimeoutError: If the code block execution exceeds the time limit
+    """
+
+    def timeout_handler(signum, frame):
+        raise TimeoutError(f"Test execution exceeded {seconds} seconds timeout")
+
+    # Set the signal handler and alarm
+    old_handler = signal.signal(signal.SIGALRM, timeout_handler)
+    signal.alarm(seconds)
+    try:
+        yield
+    finally:
+        # Restore the old signal handler and cancel the alarm
+        signal.alarm(0)
+        signal.signal(signal.SIGALRM, old_handler)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def prewarm_flashinfer_jit():
+    """Pre-warm FlashInfer JIT kernels before multi-GPU tests.
+
+    This prevents a race condition where multiple MPI ranks try to JIT-compile
+    FlashInfer kernels simultaneously to the same cache directory, causing
+    Ninja build failures like: "ninja: error: opening build log: No such file or directory"
+
+    By triggering the compilation in the main process first, the kernels are
+    cached and available for all worker ranks.
+    """
+    try:
+        import flashinfer
+        import flashinfer.page
+        import flashinfer.sampling
+
+        if torch.cuda.is_available():
+            # Prevent concurrent JIT warmup across multiple pytest processes (e.g., xdist).
+            try:
+                import fcntl  # Linux-only
+            except ImportError:
+                fcntl = None
+
+            lock_f = None
+            if fcntl is not None:
+                import pathlib
+                import tempfile
+
+                lock_path = pathlib.Path(tempfile.gettempdir()) / "flashinfer_jit_prewarm.lock"
+                lock_f = open(lock_path, "w")
+                fcntl.flock(lock_f.fileno(), fcntl.LOCK_EX)
+            # Create dummy tensors to trigger kernel JIT compilation
+            with torch.no_grad():
+                device = torch.device("cuda:0")
+
+                # Trigger page kernel compilation
+                try:
+                    # Force module loading (this triggers JIT compilation)
+                    _ = flashinfer.page.gen_page_module()
+                except Exception as exc:  # noqa: BLE001
+                    import warnings
+
+                    warnings.warn(f"FlashInfer page-kernel prewarm failed: {exc!r}", RuntimeWarning)
+
+                # Trigger sampling kernel compilation
+                try:
+                    dummy_probs = torch.softmax(torch.randn(1, 100, device=device), dim=-1)
+                    _ = flashinfer.sampling.sampling_from_probs(dummy_probs, deterministic=True)
+                except Exception as exc:  # noqa: BLE001
+                    import warnings
+
+                    warnings.warn(
+                        f"FlashInfer sampling-kernel prewarm failed: {exc!r}", RuntimeWarning
+                    )
+
+                torch.cuda.empty_cache()
+            if lock_f is not None:
+                lock_f.close()
+
+    except ImportError:
+        pass  # FlashInfer not available
+
+    yield
+
+
+@pytest.fixture(scope="module")
+def shared_dataset(llm_root):  # noqa: F811
+    """Prepare dataset once for all tests in this module."""
+    model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+    config = get_small_model_config(model_name)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        dataset_path = _prepare_dataset(
+            llm_root, temp_dir, config["args"]["model"], num_requests=10
+        )
+        # Read dataset content to return it (temp_dir will be deleted)
+        with open(dataset_path, "r") as f:
+            dataset_content = f.read()
+        yield dataset_content
+
+
+def _prepare_dataset(root_dir: str, temp_dir: str, model_path_or_name: str, num_requests: int = 10):
+    """Prepare a synthetic dataset for benchmarking."""
+    _DATASET_NAME = "synthetic_128_128.txt"
+    dataset_path = Path(temp_dir, _DATASET_NAME)
+    dataset_tool = Path(root_dir, "benchmarks", "cpp", "prepare_dataset.py")
+    script_dir = Path(root_dir, "benchmarks", "cpp")
+
+    # Generate a small dataset to run a test - matching workload configuration
+    command = [
+        "python3",
+        f"{dataset_tool}",
+        "--stdout",
+        "--tokenizer",
+        model_path_or_name,
+        "token-norm-dist",
+        "--input-mean",
+        "128",
+        "--output-mean",
+        "128",
+        "--input-stdev",
+        "0",
+        "--output-stdev",
+        "0",
+        "--num-requests",
+        str(num_requests),
+    ]
+    print(f"Running command: {' '.join(command)}")
+    result = subprocess.run(
+        command, cwd=str(script_dir), capture_output=True, text=True, timeout=300
+    )
+    if result.returncode != 0:
+        raise RuntimeError(f"Failed to prepare dataset: {result.stderr}")
+    # Grab the stdout and write it to a dataset file for passing to suite.
+    with open(dataset_path, "w") as dataset:
+        dataset.write(result.stdout)
+    return dataset_path
+
+
+@pytest.mark.parametrize(
+    "allreduce_strategy",
+    [
+        "AUTO",
+        "ONESHOT",
+        "TWOSHOT",
+        "MIN_LATENCY",
+        "NCCL",
+        "SYMM_MEM",
+    ],
+)
+def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy):  # noqa: F811
+    """Test different allreduce strategies with multi-GPU configuration making sure that there are no crashes or hangs.
+
+    Configuration:
+        The allreduce_strategy is set in the transforms config:
+        ```yaml
+        transforms:
+          detect_sharding:
+            allreduce_strategy: "ONESHOT"  # or AUTO, NCCL, TWOSHOT, etc.
+        ```
+
+    Test configuration:
+        - Model: Llama-3.1-8B with TP=2
+        - Dataset: 10 synthetic requests (128 input, 128 output tokens)
+        - Timeout: 300 seconds to catch hangs
+        - Skipped if fewer than 2 GPUs available
+
+    Args:
+        llm_root: Root directory fixture
+        shared_dataset: Shared dataset fixture (prepared once for all test runs)
+        allreduce_strategy: Strategy to test (AUTO, ONESHOT, TWOSHOT, MIN_LATENCY, NCCL)
+    """
+    # Fixed timeout for all strategies (5 minutes should be enough)
+    TEST_TIMEOUT_SECONDS = 300
+
+    model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+    config = get_small_model_config(model_name)
+    tp_size = 2
+    max_batch_size = 256
+    max_num_tokens = 8192
+
+    if not torch.cuda.is_available() or torch.cuda.device_count() < tp_size:
+        pytest.skip(f"Allreduce strategy test requires at least {tp_size} GPUs, skipping")
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Write shared dataset to temp location
+        dataset_path = Path(temp_dir, "synthetic_128_128.txt")
+        with open(dataset_path, "w") as f:
+            f.write(shared_dataset)
+
+        # Create configuration with specified allreduce strategy in transforms
+        extra_llm_api_options_path = f"{temp_dir}/extra_llm_api_options.yaml"
+        with open(extra_llm_api_options_path, "w") as f:
+            yaml.dump(
+                {
+                    **config["args"],
+                    "max_batch_size": max_batch_size,
+                    "max_num_tokens": max_num_tokens,
+                    "max_seq_len": 256,
+                    "transforms": {
+                        "detect_sharding": {
+                            "stage": "sharding",
+                            "allreduce_strategy": allreduce_strategy,
+                        },
+                        "compile_model": {
+                            "stage": "compile",
+                            "backend": "torch-cudagraph",
+                            "cuda_graph_batch_sizes": [1, 2, 4, 8, 16, 32, 64, 128, 256],
+                        },
+                    },
+                },
+                f,
+            )
+
+        # Run benchmark with specified allreduce strategy with timeout protection
+        runner = CliRunner()
+        args = [
+            "--model",
+            model_name,
+        ]
+
+        # Only pass --model_path if it's a local filesystem path
+        # Note: --model_path must come BEFORE the subcommand (throughput)
+        if str(config["args"]["model"]).startswith("/"):
+            args.extend(["--model_path", str(config["args"]["model"])])
+
+        # Add the subcommand and its options
+        args.extend(
+            [
+                "throughput",
+                "--backend",
+                "_autodeploy",
+                "--dataset",
+                str(dataset_path),
+                "--extra_llm_api_options",
+                extra_llm_api_options_path,
+                "--tp",
+                str(tp_size),
+                "--max_batch_size",
+                str(max_batch_size),
+                "--max_num_tokens",
+                str(max_num_tokens),
+            ]
+        )
+
+        try:
+            with timeout(TEST_TIMEOUT_SECONDS):
+                result = runner.invoke(main, args, catch_exceptions=False)
+                assert result.exit_code == 0, f"Benchmark failed with output: {result.output}"
+        except TimeoutError as e:
+            pytest.fail(
+                f"Test timed out after {TEST_TIMEOUT_SECONDS}s for strategy {allreduce_strategy}. "
+                f"This might indicate a hang (e.g., TWOSHOT without C++ fix). Error: {e}"
+            )

From d5163b1cbba474e2c590bdea7bf371e762fea77c Mon Sep 17 00:00:00 2001
From: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Date: Mon, 13 Jul 2026 02:01:40 -0700
Subject: [PATCH 7/8] [https://nvbugs/6403920][fix] avoid unnecessary
 block-wide barrier in single-warp RMSNorm reduction

blockReduceSum() unconditionally added a __syncthreads() plus shared-memory
cross-warp step, even for single-warp blocks (blockDim.x <= 32), where the
preceding shuffle-based warpReduceSum() already reconverges every
participating lane. For the ONESHOT strategy, whose Lamport all-reduce relies
on each thread independently busy-waiting on its own peer buffer, that
unnecessary barrier turns per-token progress into a block-wide synchronization
point and inflates tail latency under contention, observed as a ~300s CI hang
for hidden_size=64/ONESHOT (the geometry used by the AutoDeploy multigpu
smoke test's tiny Llama config) that then poisoned the following TWOSHOT case.

Skip the shared-memory/__syncthreads() path entirely when the block fits in a
single warp, matching the original validated fix's behavior for that case,
while keeping the broader multi-warp partial-block coverage.

Also add a ONESHOT variant of the existing hidden_size=64 native regression
case, mirroring the TWOSHOT one, since the gap in that coverage is what let
this regression through.

Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
---
 .../allReduceFusionKernels.cu                  | 18 +++++++++++++++---
 .../kernels/allReduce/allReduceKernelTest.cu   |  8 +++++++-
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
index 397bc0032b3f..116da185b1d5 100644
--- a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
+++ b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
@@ -49,7 +49,6 @@ __device__ __forceinline__ float warpReduceSum(float val, uint32_t mask, int act
 
 __device__ __forceinline__ void blockReduceSum(float* val)
 {
-    __shared__ float shared[32];
     int const lane = threadIdx.x & (warpSize - 1);
     int const warpId = threadIdx.x / warpSize;
     int const warpStart = warpId * warpSize;
@@ -58,6 +57,20 @@ __device__ __forceinline__ void blockReduceSum(float* val)
     uint32_t const mask = getLaneMask(activeLanes);
 
     *val = warpReduceSum(*val, mask, activeLanes);
+
+    int const warpCount = (blockDim.x + warpSize - 1) / warpSize;
+    if (warpCount == 1)
+    {
+        // Single-warp block: the shuffle-based warp reduction above already reconverges every
+        // participating lane, so skip the block-wide barrier below. It matters here because this
+        // path also serves the oneshot Lamport all-reduce, whose per-thread busy-wait polling
+        // relies on threads making independent progress; an unconditional __syncthreads() would
+        // turn that into a per-token, block-wide synchronization point and inflate tail latency
+        // under contention (observed as a ~300s CI timeout for hidden_size=64/ONESHOT).
+        return;
+    }
+
+    __shared__ float shared[32];
     if (lane == 0)
     {
         shared[warpId] = *val;
@@ -65,8 +78,7 @@ __device__ __forceinline__ void blockReduceSum(float* val)
 
     __syncthreads();
 
-    int const warpCount = (blockDim.x + warpSize - 1) / warpSize;
-    if (warpCount > 1 && warpId == 0)
+    if (warpId == 0)
     {
         *val = lane < warpCount ? shared[lane] : 0.f;
         *val = warpReduceSum(*val, 0xffffffffU, warpSize);
diff --git a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu
index f2e01e6c5048..87014f41b430 100644
--- a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu
+++ b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu
@@ -623,8 +623,14 @@ TEST(Kernel, AllReduce)
 
     // Exercise the partial-warp RMSNorm reduction used by small AutoDeploy test models.
     // At hidden_size=64, a 16-bit dtype processes eight 128-bit accesses per token, so the
-    // all-reduce kernel launches fewer than one warp of threads.
+    // all-reduce kernel launches fewer than one warp of threads. Cover both ONESHOT and TWOSHOT:
+    // they share this reduction code but drive it through different communication paths (ONESHOT's
+    // Lamport buffers rely on per-thread busy-wait polling instead of TWOSHOT's barrier), so a
+    // regression can affect one strategy without the other.
     Workspace small_hidden_workspace(world_size, rank, 8192, 64);
+    pass = pass
+        && test(small_hidden_workspace, 8192, 64, false, true, 3, 3, AllReduceStrategyType::ONESHOT,
+            AllReduceStrategyConfig(0), AllReduceFusionOp::RESIDUAL_RMS_NORM);
     pass = pass
         && test(small_hidden_workspace, 8192, 64, false, true, 3, 3, AllReduceStrategyType::TWOSHOT,
             AllReduceStrategyConfig(0), AllReduceFusionOp::RESIDUAL_RMS_NORM);

From 3c7abf943ac6c167ed2e294d76bc661a4b15a08f Mon Sep 17 00:00:00 2001
From: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Date: Mon, 13 Jul 2026 08:31:36 -0700
Subject: [PATCH 8/8] [None][test] Probe: use full-warp hidden_size in
 allreduce smoke test

Diagnostic-only change: overrides the tiny test model's hidden_size from
64 (which launches an 8-thread, partial-warp block in the fused all-reduce
RMSNorm kernel) to 256 (exactly one full warp), so the smoke test no
longer exercises the partial-warp reduction path at all.

Two prior CI runs of the base branch (agent/fix-ad-allreduce-ci) failed
with an identical CUBLAS_STATUS_EXECUTION_FAILED error in an unrelated
cuBLAS GEMM (torch_linear_simple), on the same node
(nsc-svg-slurm-1-gpu-120), at a different allreduce strategy each time.
This branch isolates whether that failure is specific to the partial-warp
kernel path or is unrelated node/infra flakiness that also reproduces on
the full-warp path.

Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
---
 .../smoke/test_ad_allreduce_strategies.py         | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py
index 937bdd4c3ff9..48b650c3ba1f 100644
--- a/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py
+++ b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py
@@ -217,7 +217,20 @@ def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy):  #
     TEST_TIMEOUT_SECONDS = 300
 
     model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-    config = get_small_model_config(model_name)
+    # Override hidden_size to a multiple of one warp's worth of 128-bit accesses (32 threads *
+    # 8 fp16 elements/access = 256) so the fused all-reduce/RMSNorm kernel never launches a
+    # partial-warp block. This isolates whether CI failures are specific to the partial-warp
+    # code path or are unrelated infra flakiness that also affects the full-warp path.
+    config = get_small_model_config(
+        model_name,
+        model_kwargs={
+            "num_hidden_layers": 1,
+            "hidden_size": 256,
+            "intermediate_size": 256,
+            "num_attention_heads": 2,
+            "num_key_value_heads": 1,
+        },
+    )
     tp_size = 2
     max_batch_size = 256
     max_num_tokens = 8192