diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu index 5a3edda04a70..116da185b1d5 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,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(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 struct SyncComm { @@ -296,7 +357,14 @@ protected: float v = static_cast(reinterpret_cast(&residual)[i]); acc += v * v; } - tensorrt_llm::common::blockReduceSumV2(&acc); + if (blockDim.x % warpSize != 0) + { + blockReduceSum(&acc); + } + else + { + tensorrt_llm::common::blockReduceSumV2(&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/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 == 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(), ref_fp8_output.device_data<__nv_fp8_e4m3>(), message_size, + quantize_to_fp8(ref_output.device_data(), ref_fp8_output.device_data<__nv_fp8_e4m3>(), message_size, m_scale_factor.device_data(), 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 } 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..87014f41b430 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,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) { diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 4917713b5dda..04452b35804a 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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) 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..48b650c3ba1f --- /dev/null +++ b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py @@ -0,0 +1,310 @@ +# 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" + # 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 + + 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}" + )