[Common] CGEMM stream ID logging

docs/envvars.rst

@@ -460,6 +460,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)

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" \

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)
@@ -95,11 +100,25 @@ 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})
+
+# 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)
 
 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).

tests/cpp_distributed/test_comm_gemm_stream_logging.cu

@@ -0,0 +1,68 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+#include <mpi.h>
+#include <nccl.h>
+#include <transformer_engine/comm_gemm.h>
+
+#include <cstdlib>
+#include <string>
+
+#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));
+}

transformer_engine/common/comm_gemm/comm_gemm.cpp

@@ -9,6 +9,7 @@
 #include <cuda_runtime.h>
 #include <transformer_engine/recipe.h>
 
+#include <iostream>
 #include <map>
 #include <memory>
 #include <mutex>
@@ -21,6 +22,7 @@
 
 #include "../common.h"
 #include "../util/logging.h"
+#include "../util/system.h"
 
 #ifdef NVTE_WITH_CUBLASMP
 
@@ -49,6 +51,20 @@ CudaStream CudaStreamCreate() {
   return CreateWithCudaCheck<cudaStream_t>(cudaStreamCreate, cudaStreamDestroy);
 }
 
+bool LogCollectiveGemmStreamIds() {
+  static const bool enabled =
+      transformer_engine::getenv<bool>("NVTE_LOG_COLLECTIVE_GEMM_STREAM_IDS", false);
+  return enabled;
+}
+
+void LogCollectiveGemmStreamId(cudaStream_t stream, int rank) {
+  if (!LogCollectiveGemmStreamIds()) return;
+
+  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;
+}
+
 using CudaEvent = std::unique_ptr<std::remove_pointer_t<cudaEvent_t>, decltype(&cudaEventDestroy)>;
 
 CudaEvent CudaEventCreate(unsigned flags) {
@@ -633,6 +649,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());