From dc13d311115adcb200ba1d50d30c8ead5f9bf5eb Mon Sep 17 00:00:00 2001
From: chen <cx2016013@163.com>
Date: Thu, 21 May 2026 09:57:38 +0000
Subject: [PATCH] fix: correct batched matmul strides for bs=1 and integrate
 CTest into build pipeline

Set stride to 0 when batch_size is 1 to enable proper broadcasting in cuBLAS,
and add configurable CTest execution after builds with googletest submodule.
---
 docs/test_usage_guide.md                |  2 +-
 infini_train/src/kernels/cuda/matmul.cu | 18 +++++++++---------
 scripts/run_models_and_profile.bash     |  5 +++++
 scripts/test_config.json                |  4 +++-
 third_party/googletest                  |  1 +
 5 files changed, 19 insertions(+), 11 deletions(-)
 create mode 160000 third_party/googletest

diff --git a/docs/test_usage_guide.md b/docs/test_usage_guide.md
index 108f995a..45113e52 100644
--- a/docs/test_usage_guide.md
+++ b/docs/test_usage_guide.md
@@ -8,7 +8,7 @@
 
 ```bash
 mkdir build && cd build
-cmake -DBUILD_TEST=ON -DUSE_CUDA=ON ..
+cmake -DBUILD_TEST=ON -DUSE_CUDA=ON -DUSE_NCCL=ON ..
 make -j$(nproc)
 ```
 
diff --git a/infini_train/src/kernels/cuda/matmul.cu b/infini_train/src/kernels/cuda/matmul.cu
index 8ee33174..7e301039 100644
--- a/infini_train/src/kernels/cuda/matmul.cu
+++ b/infini_train/src/kernels/cuda/matmul.cu
@@ -63,9 +63,9 @@ std::shared_ptr<Tensor> MatmulForward(const std::shared_ptr<Tensor> &input, cons
                          .alpha = 1.0f,
                          .beta = 0.0f,
                          .batch_count = static_cast<int>(bs),
-                         .stride_a = n * k,
-                         .stride_b = k * m,
-                         .stride_c = m * n,
+                         .stride_a = bs > 1 ? n * k : 0,
+                         .stride_b = bs > 1 ? k * m : 0,
+                         .stride_c = bs > 1 ? m * n : 0,
                          .input_dtype = dtype,
                          .output_dtype = dtype,
                      });
@@ -133,9 +133,9 @@ std::shared_ptr<Tensor> MatmulBackwardInput(const std::shared_ptr<Tensor> &other
                          .alpha = 1.0f,
                          .beta = 0.0f,
                          .batch_count = static_cast<int>(bs),
-                         .stride_a = k * n,
-                         .stride_b = n * m,
-                         .stride_c = m * k,
+                         .stride_a = bs > 1 ? k * n : 0,
+                         .stride_b = bs > 1 ? n * m : 0,
+                         .stride_c = bs > 1 ? m * k : 0,
                          .input_dtype = compute_dtype,
                          .output_dtype = output_dtype,
                      });
@@ -202,9 +202,9 @@ std::shared_ptr<Tensor> MatmulBackwardOther(const std::shared_ptr<Tensor> &input
                          .alpha = 1.0f,
                          .beta = 0.0f,
                          .batch_count = static_cast<int>(bs),
-                         .stride_a = n * m,
-                         .stride_b = k * m,
-                         .stride_c = n * k,
+                         .stride_a = bs > 1 ? n * m : 0,
+                         .stride_b = bs > 1 ? k * m : 0,
+                         .stride_c = bs > 1 ? n * k : 0,
                          .input_dtype = compute_dtype,
                          .output_dtype = output_dtype,
                      });
diff --git a/scripts/run_models_and_profile.bash b/scripts/run_models_and_profile.bash
index 06589904..8d68a08a 100755
--- a/scripts/run_models_and_profile.bash
+++ b/scripts/run_models_and_profile.bash
@@ -70,6 +70,8 @@ BUILD_DIR="$(read_var BUILD_DIR)";              : "${BUILD_DIR:=../build}"
 LOG_DIR="$(read_var LOG_DIR)";                  : "${LOG_DIR:=logs}"
 PROFILE_LOG_DIR="$(read_var PROFILE_LOG_DIR)";  : "${PROFILE_LOG_DIR:=./profile_logs}"
 COMPARE_LOG_DIR="$(read_var COMPARE_LOG_DIR)";  : "${COMPARE_LOG_DIR:=}"
+RUN_CTEST="$(read_var RUN_CTEST)";              : "${RUN_CTEST:=true}"
+CTEST_CMD="$(read_var CTEST_CMD)";              : "${CTEST_CMD:=ctest --output-on-failure -LE cuda -j$(nproc) && ctest --output-on-failure -L cuda -j1}"
 
 mkdir -p "$BUILD_DIR" "$LOG_DIR" "$PROFILE_LOG_DIR"
 
@@ -244,6 +246,9 @@ for ((id=0; id<num_builds; ++id)); do
     # always clean before another build
     clean_build_dir
     run_and_log "$LAST_CMAKE_CMD" "${build_id}" "no" "build"
+    if [[ "$RUN_CTEST" == "true" ]]; then
+        run_and_log "$CTEST_CMD" "ctest_${build_id}" "no" "ctest"
+    fi
 
     # profile flag for runs
     profile_flag="no"
diff --git a/scripts/test_config.json b/scripts/test_config.json
index 54332f70..1700f897 100644
--- a/scripts/test_config.json
+++ b/scripts/test_config.json
@@ -7,7 +7,9 @@
         "LLAMA3_LLMC_FILEPATH": "/data/shared/InfiniTrain-dev/data/llmc/llama3/llama3.2_1B_fp32.bin",
         "PROFILE_LOG_DIR": "./profile_logs",
         "LOG_DIR": "./logs",
-        "COMPARE_LOG_DIR": ""
+        "COMPARE_LOG_DIR": "",
+        "RUN_CTEST": "true",
+        "CTEST_CMD": "ctest --output-on-failure -LE cuda -j$(nproc) && ctest --output-on-failure -L cuda -j1"
     },
     "builds": [
         {
diff --git a/third_party/googletest b/third_party/googletest
new file mode 160000
index 00000000..f8d7d77c
--- /dev/null
+++ b/third_party/googletest
@@ -0,0 +1 @@
+Subproject commit f8d7d77c06936315286eb55f8de22cd23c188571