From e35b4bc974e570854f08db3fa27a9c0d6b46895d Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 22 Oct 2026 09:13:47 -0700 Subject: [PATCH 1/5] CGEMM Stream ID logging Signed-off-by: Jeremy Berchtold --- tests/cpp_distributed/test_comm_gemm.cu | 29 +++++++++++++++++++ .../common/comm_gemm/comm_gemm.cpp | 17 +++++++++++ 2 files changed, 46 insertions(+) diff --git a/tests/cpp_distributed/test_comm_gemm.cu b/tests/cpp_distributed/test_comm_gemm.cu index b1c4a9395a9..7e2f31362a8 100644 --- a/tests/cpp_distributed/test_comm_gemm.cu +++ b/tests/cpp_distributed/test_comm_gemm.cu @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -225,6 +226,34 @@ float GetScale(float amax) { return static_cast(static_cast(std::numeric_limits::max())) / amax; } +TEST(CommGemmStreamLogging, LogsStreamIdToStdout) { + int nranks{}; + int rank{}; + CHECK_MPI(MPI_Comm_size(MPI_COMM_WORLD, &nranks)); + CHECK_MPI(MPI_Comm_rank(MPI_COMM_WORLD, &rank)); + NVTE_CHECK_CUDA(cudaSetDevice(rank)); + + ncclUniqueId id{}; + if (rank == 0) CHECK_NCCL(ncclGetUniqueId(&id)); + CHECK_MPI(MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD)); + + ncclComm_t comm{}; + CHECK_NCCL(ncclCommInitRank(&comm, nranks, id, rank)); + ASSERT_EQ(setenv("NVTE_LOG_COLLECTIVE_GEMM_STREAM_IDS", "1", 1), 0); + + testing::internal::CaptureStdout(); + NVTECommGemmCtx* ctx = nvte_comm_gemm_ctx_create(comm, nranks, rank); + const std::string output = testing::internal::GetCapturedStdout(); + + const std::string prefix = + "NVTE collective GEMM stream ID (rank " + std::to_string(rank) + "): "; + EXPECT_NE(output.find(prefix), std::string::npos) << "Captured stdout: " << output; + + ASSERT_EQ(unsetenv("NVTE_LOG_COLLECTIVE_GEMM_STREAM_IDS"), 0); + nvte_comm_gemm_ctx_destroy(ctx); + CHECK_NCCL(ncclCommDestroy(comm)); +} + struct Params { DType a_type; DType b_type; diff --git a/transformer_engine/common/comm_gemm/comm_gemm.cpp b/transformer_engine/common/comm_gemm/comm_gemm.cpp index 3545347da47..7beecc2850f 100644 --- a/transformer_engine/common/comm_gemm/comm_gemm.cpp +++ b/transformer_engine/common/comm_gemm/comm_gemm.cpp @@ -8,6 +8,7 @@ #include +#include #include #include #include @@ -19,6 +20,7 @@ #include "../common.h" #include "../util/logging.h" +#include "../util/system.h" #ifdef NVTE_WITH_CUBLASMP @@ -47,6 +49,20 @@ CudaStream CudaStreamCreate() { return CreateWithCudaCheck(cudaStreamCreate, cudaStreamDestroy); } +bool LogCollectiveGemmStreamIds() { + static const bool enabled = + transformer_engine::getenv("NVTE_LOG_COLLECTIVE_GEMM_STREAM_IDS", false); + return enabled; +} + +void LogCollectiveGemmStreamId(cudaStream_t stream, int rank) { + if (!LogCollectiveGemmStreamIds()) return; + + unsigned long long stream_id{}; + NVTE_CHECK_CUDA(cudaStreamGetId(stream, &stream_id)); + std::cout << "NVTE collective GEMM stream ID (rank " << rank << "): " << stream_id << std::endl; +} + using CudaEvent = std::unique_ptr, decltype(&cudaEventDestroy)>; CudaEvent CudaEventCreate(unsigned flags) { @@ -481,6 +497,7 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo NVTECommGemmCtx* nvte_comm_gemm_ctx_create(ncclComm_t comm, int nranks, int rank) { NVTE_API_CALL(nvte_comm_gemm_ctx_create); auto stream = CudaStreamCreate(); + LogCollectiveGemmStreamId(stream.get(), rank); auto event = CudaEventCreate(cudaEventDisableTiming); auto cublas_mp = CublasMpCreate(stream.get()); From 4383a8367c55dfe6cb9a2f7ec07c908d97b823b0 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 22 Oct 2026 09:21:41 -0700 Subject: [PATCH 2/5] Document collective GEMM stream ID logging Signed-off-by: Jeremy Berchtold --- docs/envvars.rst | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/envvars.rst b/docs/envvars.rst index 97eaed5ddc6..2adbf183273 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -416,6 +416,15 @@ GEMM Configuration :Default: Total SM count :Description: External SM margin for GEMM operations. Specifies the number of SMs to use for GEMM operations. The actual number of SMs used is ``sm_count - NVTE_EXT_MARGIN_SM``. +.. envvar:: NVTE_LOG_COLLECTIVE_GEMM_STREAM_IDS + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Log the CUDA stream ID and distributed rank to standard output when Transformer + Engine creates the internal CUDA stream for a collective GEMM context. The value + is read and cached on first use, so changing the environment variable later does + not affect that process. + .. envvar:: NVTE_AG_P2P_MULTI_ATOMIC :Type: ``int`` (0 or 1) From 82ce4c7ffa2a04e5ae3bdc06303f8340054e2a2c Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 22 Oct 2026 10:06:26 -0700 Subject: [PATCH 3/5] Isolate collective GEMM stream logging test Signed-off-by: Jeremy Berchtold --- qa/L1_cpp_distributed/test.sh | 10 +++ tests/cpp_distributed/CMakeLists.txt | 10 +++ tests/cpp_distributed/test_comm_gemm.cu | 29 -------- .../test_comm_gemm_stream_logging.cu | 68 +++++++++++++++++++ .../common/comm_gemm/comm_gemm.cpp | 2 +- 5 files changed, 89 insertions(+), 30 deletions(-) create mode 100644 tests/cpp_distributed/test_comm_gemm_stream_logging.cu diff --git a/qa/L1_cpp_distributed/test.sh b/qa/L1_cpp_distributed/test.sh index 4122163f695..555c8b08882 100755 --- a/qa/L1_cpp_distributed/test.sh +++ b/qa/L1_cpp_distributed/test.sh @@ -28,8 +28,11 @@ if [[ $(nvidia-smi --list-gpus | wc -l) -ge 4 ]]; then # mask the other's results. Skip mpirun when the binary is missing. if [[ $configure_ok -eq 1 ]]; then comm_gemm_ok=1 + comm_gemm_stream_logging_ok=1 ep_ok=1 cmake --build build --target test_comm_gemm || { test_fail "test_comm_gemm_build"; comm_gemm_ok=0; } + cmake --build build --target test_comm_gemm_stream_logging \ + || { test_fail "test_comm_gemm_stream_logging_build"; comm_gemm_stream_logging_ok=0; } cmake --build build --target test_ep || { test_fail "test_ep_build"; ep_ok=0; } if [[ $comm_gemm_ok -eq 1 ]]; then @@ -39,6 +42,13 @@ if [[ $(nvidia-smi --list-gpus | wc -l) -ge 4 ]]; then || test_fail "test_comm_gemm" fi + if [[ $comm_gemm_stream_logging_ok -eq 1 ]]; then + # Run in a separate process because the logging setting is cached on first use. + mpirun --allow-run-as-root --np 4 --oversubscribe bash -c \ + "exec ./build/test_comm_gemm_stream_logging --gtest_output=xml:$XML_LOG_DIR/cpp_distributed_test_comm_gemm_stream_logging.rank\${OMPI_COMM_WORLD_RANK}.xml" \ + || test_fail "test_comm_gemm_stream_logging" + fi + if [[ $ep_ok -eq 1 ]]; then # EP suites; runner self-skips on pre-Hopper GPUs. GTEST_XML_PREFIX="$XML_LOG_DIR/cpp_distributed_test_ep" \ diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index 13b6242816a..c75c7cc7209 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -43,6 +43,9 @@ add_executable(test_comm_gemm test_comm_gemm.cu ../cpp/test_common.cu) +add_executable(test_comm_gemm_stream_logging + test_comm_gemm_stream_logging.cu) + list(APPEND test_comm_gemm_INCLUDES ${CMAKE_SOURCE_DIR}/../../transformer_engine/common/include ${CMAKE_SOURCE_DIR}/../../transformer_engine/common @@ -51,6 +54,8 @@ list(APPEND test_comm_gemm_INCLUDES ${MPI_CXX_INCLUDE_PATH} $ENV{CUBLASMP_HOME}/include) target_include_directories(test_comm_gemm PRIVATE ${test_comm_gemm_INCLUDES}) +target_include_directories(test_comm_gemm_stream_logging + PRIVATE ${test_comm_gemm_INCLUDES}) find_package(CUDAToolkit REQUIRED) find_package(OpenMP REQUIRED) @@ -91,11 +96,16 @@ list(APPEND test_comm_gemm_LINKER_LIBS OpenMP::OpenMP_CXX MPI::MPI_CXX) target_link_libraries(test_comm_gemm PUBLIC ${test_comm_gemm_LINKER_LIBS}) +target_link_libraries(test_comm_gemm_stream_logging + PUBLIC ${test_comm_gemm_LINKER_LIBS}) target_compile_options(test_comm_gemm PRIVATE -O2 -fopenmp) +target_compile_options(test_comm_gemm_stream_logging PRIVATE -O2 -fopenmp) include(GoogleTest) gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600) +# This test must run in its own MPI-launched process because the environment +# variable it exercises is cached for the lifetime of the process. # -- EP distributed tests ------------------------------------------------------ # Launched via mpirun; ncclUniqueId exchange uses MPI_Bcast (see test_ep_common.h). diff --git a/tests/cpp_distributed/test_comm_gemm.cu b/tests/cpp_distributed/test_comm_gemm.cu index 7e2f31362a8..b1c4a9395a9 100644 --- a/tests/cpp_distributed/test_comm_gemm.cu +++ b/tests/cpp_distributed/test_comm_gemm.cu @@ -13,7 +13,6 @@ #include #include -#include #include #include #include @@ -226,34 +225,6 @@ float GetScale(float amax) { return static_cast(static_cast(std::numeric_limits::max())) / amax; } -TEST(CommGemmStreamLogging, LogsStreamIdToStdout) { - int nranks{}; - int rank{}; - CHECK_MPI(MPI_Comm_size(MPI_COMM_WORLD, &nranks)); - CHECK_MPI(MPI_Comm_rank(MPI_COMM_WORLD, &rank)); - NVTE_CHECK_CUDA(cudaSetDevice(rank)); - - ncclUniqueId id{}; - if (rank == 0) CHECK_NCCL(ncclGetUniqueId(&id)); - CHECK_MPI(MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD)); - - ncclComm_t comm{}; - CHECK_NCCL(ncclCommInitRank(&comm, nranks, id, rank)); - ASSERT_EQ(setenv("NVTE_LOG_COLLECTIVE_GEMM_STREAM_IDS", "1", 1), 0); - - testing::internal::CaptureStdout(); - NVTECommGemmCtx* ctx = nvte_comm_gemm_ctx_create(comm, nranks, rank); - const std::string output = testing::internal::GetCapturedStdout(); - - const std::string prefix = - "NVTE collective GEMM stream ID (rank " + std::to_string(rank) + "): "; - EXPECT_NE(output.find(prefix), std::string::npos) << "Captured stdout: " << output; - - ASSERT_EQ(unsetenv("NVTE_LOG_COLLECTIVE_GEMM_STREAM_IDS"), 0); - nvte_comm_gemm_ctx_destroy(ctx); - CHECK_NCCL(ncclCommDestroy(comm)); -} - struct Params { DType a_type; DType b_type; diff --git a/tests/cpp_distributed/test_comm_gemm_stream_logging.cu b/tests/cpp_distributed/test_comm_gemm_stream_logging.cu new file mode 100644 index 00000000000..2bcec8def38 --- /dev/null +++ b/tests/cpp_distributed/test_comm_gemm_stream_logging.cu @@ -0,0 +1,68 @@ +/************************************************************************* + * Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include + +#include +#include + +#define CHECK_MPI(expr) \ + do { \ + int err = (expr); \ + if (err != MPI_SUCCESS) { \ + char err_str[MPI_MAX_ERROR_STRING + 1]{}; \ + int len{}; \ + MPI_Error_string(err, err_str, &len); \ + EXPECT_TRUE(false) << "MPI error: " << err << ": " << err_str; \ + } \ + } while (false) + +#define CHECK_NCCL(expr) \ + do { \ + ncclResult_t err = (expr); \ + if (err != ncclSuccess) { \ + EXPECT_TRUE(false) << "NCCL error: " << err << ": " << ncclGetErrorString(err); \ + } \ + } while (false) + +int main(int argc, char* argv[]) { + ::testing::InitGoogleTest(&argc, argv); + CHECK_MPI(MPI_Init(&argc, &argv)); + auto ret = RUN_ALL_TESTS(); + CHECK_MPI(MPI_Finalize()); + return ret; +} + +TEST(CommGemmStreamLogging, LogsStreamIdToStdout) { + int nranks{}; + int rank{}; + CHECK_MPI(MPI_Comm_size(MPI_COMM_WORLD, &nranks)); + CHECK_MPI(MPI_Comm_rank(MPI_COMM_WORLD, &rank)); + ASSERT_EQ(cudaSetDevice(rank), cudaSuccess); + + ncclUniqueId id{}; + if (rank == 0) CHECK_NCCL(ncclGetUniqueId(&id)); + CHECK_MPI(MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD)); + + ncclComm_t comm{}; + CHECK_NCCL(ncclCommInitRank(&comm, nranks, id, rank)); + ASSERT_EQ(setenv("NVTE_LOG_COLLECTIVE_GEMM_STREAM_IDS", "1", 1), 0); + + testing::internal::CaptureStdout(); + NVTECommGemmCtx* ctx = nvte_comm_gemm_ctx_create(comm, nranks, rank); + const std::string output = testing::internal::GetCapturedStdout(); + + const std::string prefix = + "NVTE collective GEMM stream ID (rank " + std::to_string(rank) + "): "; + EXPECT_NE(output.find(prefix), std::string::npos) << "Captured stdout: " << output; + + nvte_comm_gemm_ctx_destroy(ctx); + CHECK_NCCL(ncclCommDestroy(comm)); +} diff --git a/transformer_engine/common/comm_gemm/comm_gemm.cpp b/transformer_engine/common/comm_gemm/comm_gemm.cpp index 7beecc2850f..7d2b4ea94c0 100644 --- a/transformer_engine/common/comm_gemm/comm_gemm.cpp +++ b/transformer_engine/common/comm_gemm/comm_gemm.cpp @@ -58,7 +58,7 @@ bool LogCollectiveGemmStreamIds() { void LogCollectiveGemmStreamId(cudaStream_t stream, int rank) { if (!LogCollectiveGemmStreamIds()) return; - unsigned long long stream_id{}; + unsigned long long stream_id{}; // NOLINT(runtime/int) NVTE_CHECK_CUDA(cudaStreamGetId(stream, &stream_id)); std::cout << "NVTE collective GEMM stream ID (rank " << rank << "): " << stream_id << std::endl; } From fe9222e35c11c245b65cfd502b4be3da901d9f88 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Thu, 22 Oct 2026 10:17:49 -0700 Subject: [PATCH 4/5] Fix collective GEMM test copyright header Signed-off-by: Jeremy Berchtold --- tests/cpp_distributed/test_comm_gemm_stream_logging.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cpp_distributed/test_comm_gemm_stream_logging.cu b/tests/cpp_distributed/test_comm_gemm_stream_logging.cu index 2bcec8def38..4b782ef6894 100644 --- a/tests/cpp_distributed/test_comm_gemm_stream_logging.cu +++ b/tests/cpp_distributed/test_comm_gemm_stream_logging.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ From 63225a21d6a346b8e25ee219e15199f2c6a8eadf Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Fri, 23 Oct 2026 11:20:05 -0700 Subject: [PATCH 5/5] Link collective GEMM stream logging test against Python Signed-off-by: Jeremy Berchtold --- tests/cpp_distributed/CMakeLists.txt | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index c75c7cc7209..cae373a4e48 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -99,6 +99,15 @@ target_link_libraries(test_comm_gemm PUBLIC ${test_comm_gemm_LINKER_LIBS}) target_link_libraries(test_comm_gemm_stream_logging PUBLIC ${test_comm_gemm_LINKER_LIBS}) +# CuTeDSL-enabled core libraries require standalone executables to provide CPython symbols. +if(NVTE_WITH_CUTEDSL OR "$ENV{NVTE_WITH_CUTEDSL}") + find_package(Python COMPONENTS Interpreter Development.Embed REQUIRED) + target_link_libraries(test_comm_gemm_stream_logging PRIVATE + "-Wl,--no-as-needed" + Python::Python + "-Wl,--as-needed") +endif() + target_compile_options(test_comm_gemm PRIVATE -O2 -fopenmp) target_compile_options(test_comm_gemm_stream_logging PRIVATE -O2 -fopenmp)