[None][test] Probe: full-warp hidden_size in allreduce CI smoke test

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,67 @@ TRTLLM_NAMESPACE_BEGIN
 
 namespace kernels::ar_fusion
 {
+__device__ __forceinline__ uint32_t getLaneMask(int activeLanes)
+{
+    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)
+    {
+        if (offset < activeLanes)
+        {
+            float const other = __shfl_xor_sync(mask, val, offset, warpSize);
+            if ((lane ^ offset) < activeLanes)
+            {
+                val += other;
+            }
+        }
+    }
+    return val;
+}
+
+__device__ __forceinline__ void blockReduceSum(float* val)
+{
+    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);
+
+    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;
+    }
+
+    __syncthreads();
+
+    if (warpId == 0)
+    {
+        *val = lane < warpCount ? shared[lane] : 0.f;
+        *val = warpReduceSum(*val, 0xffffffffU, warpSize);
+    }
+}
+
 template <int NRanks>
 struct SyncComm
 {
@@ -296,7 +357,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 != 0)
+        {
+            blockReduceSum(&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)

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
 }
 

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,22 @@ 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. 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);
+    EXPECT_TRUE(pass);
+
     ops[0] = AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM;
     for (auto config : configs)
     {

tests/integration/test_lists/waives.txt

@@ -235,9 +235,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_B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] SKIP (https://nvbugs/6432831)
 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)
@@ -443,10 +440,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_kv_transfer.py::test_transfer_worker_v2[tp4_pp1_to_tp2_pp2] SKIP (https://nvbugs/6426834)