diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index fe95722ab5dc..06332fc50407 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -563,19 +563,29 @@ class AllreduceOp inputPtr = windowBuffer0.ptr; } - // Use window-backed output buffer - auto [normOut, windowBuffer1] = createNCCLWindowTensor(rawComm, input.sizes(), input.scalar_type()); - torch::Tensor outputTensor = windowBuffer1.isValid() ? normOut : torch::empty_like(inputTensor); - void* outputPtr = windowBuffer1.isValid() ? windowBuffer1.ptr : outputTensor.data_ptr(); - if (!windowBuffer1.isValid()) + // Request an output window only when the message meets the same + // registration threshold used by the input path. Smaller messages use a + // regular output tensor and skip the window allocation path. + torch::Tensor outputTensor; + if (bufferSizeBytes >= minRegistrationThreshold) + { + auto [windowOutput, windowBuffer1] = createNCCLWindowTensor(rawComm, input.sizes(), input.scalar_type()); + if (windowBuffer1.isValid()) + { + outputTensor = windowOutput; + } + } + if (!outputTensor.defined()) { TLLM_LOG_DEBUG( "[runNCCLAllReduceSymmetric] No valid symmetric buffer available; " "using plain CUDA tensor for output"); + outputTensor = torch::empty_like(inputTensor); } // Perform allreduce - NCCLCHECK_THROW(ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, comm, stream)); + NCCLCHECK_THROW( + ncclAllReduce(inputPtr, outputTensor.data_ptr(), size, (*getDtypeMap())[mType], ncclSum, comm, stream)); if (mOp == AllReduceFusionOp::NONE) { diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 3911a2e9151d..37c5e2efe419 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -6193,6 +6193,7 @@ def launchTestJobs(pipeline, testFilter, globalVars) "DGX_A100-FMHA-Post-Merge-1": ["auto:dgx-a100-x1", "l0_a100", 1, 1], "DGX_H100-2_GPUs-PyTorch-Others-1": ["auto:dgx-h100-x2", "l0_dgx_h100", 1, 2, 2], "DGX_H100-2_GPUs-PyTorch-Others-2": ["auto:dgx-h100-x2", "l0_dgx_h100", 2, 2, 2], + "DGX_H100-2_GPUs-PyTorch-Others-Post-Merge-1": ["auto:dgx-h100-x2", "l0_dgx_h100", 1, 1, 2], "DGX_H100-2_GPUs-PyTorch-GptOss-1": ["auto:dgx-h100-x2", "l0_dgx_h100", 1, 1, 2], "DGX_H100-2_GPUs-PyTorch-Ray-1": ["auto:dgx-h100-x2", "l0_dgx_h100", 1, 1, 2], "DGX_H100-4_GPUs-PyTorch-DeepSeek-1": ["auto:dgx-h100-x4", "l0_dgx_h100", 1, 1, 4], diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 43ab57a68f7a..83997d5341a5 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -343,3 +343,19 @@ l0_dgx_h100: - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_off-cpp_mamba_cache] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_on-python_mamba_cache] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_on-cpp_mamba_cache] +- condition: + ranges: + system_gpu_count: + gte: 2 + lte: 2 + wildcards: + gpu: + - '*h100*' + linux_distribution_name: ubuntu* + terms: + stage: post_merge + backend: pytorch + auto_trigger: others + orchestrator: mpi + tests: + - unittest/_torch/multi_gpu -m "post_merge" TIMEOUT (90) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index c3a649190f0f..72b5d5893b38 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -345,7 +345,6 @@ unittest/_torch/modeling/test_modeling_nemotron_nano_v2_vl.py::test_nemotron_nan unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124) unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py::TestQwen3LoRA::test_qwen3_fp8_lora SKIP (https://nvbugs/6668777) unittest/_torch/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] SKIP (https://nvbugs/5989912) -unittest/_torch/multi_gpu/test_linear.py::test_row_linear_norm_fusion[2-hidden:16-seqlen:2] SKIP (https://nvbugs/6501404) unittest/_torch/speculative/hw_agnostic/test_dflash.py::test_dflash_qwen3_5_4b[False] SKIP (https://nvbugs/6535767) unittest/_torch/speculative/hw_agnostic/test_dflash.py::test_dflash_qwen3_5_4b[True] SKIP (https://nvbugs/6535767) unittest/_torch/speculative/hw_agnostic/test_ngram.py::test_llama_ngram[True-True-TRTLLM] SKIP (https://nvbugs/6507102) diff --git a/tests/unittest/_torch/multi_gpu/test_linear.py b/tests/unittest/_torch/multi_gpu/test_linear.py index 4db1c6b053ce..3e4b40b2f57c 100644 --- a/tests/unittest/_torch/multi_gpu/test_linear.py +++ b/tests/unittest/_torch/multi_gpu/test_linear.py @@ -310,6 +310,7 @@ def test_row_linear(hidden_size, mpi_pool_executor): assert r is True +@pytest.mark.post_merge @pytest.mark.skipif(torch.cuda.device_count() < 2, reason='needs 2 GPUs to run this test') @pytest.mark.parametrize("seq_len", [2, 32], ids=lambda x: f"seqlen:{x}")