feat: added gtest_main with filters based on compute_capabilities

src/common/CMakeLists.txt

@@ -28,6 +28,19 @@ cc_library(
     glog::glog
 )
 
+cc_library(
+  TESTONLY
+  NAME
+    gtest_main
+  SRCS
+    gtest_main.cpp
+  DEPS
+    GTest::gtest
+    CUDA::toolkit
+  LINKOPTS
+    cudart
+)
+
 cc_test(
   NAME 
     common_test
@@ -39,7 +52,7 @@ cc_test(
     common
     absl::synchronization
     absl::time
-    GTest::gtest_main
+    :gtest_main
     torch
 )
 

src/common/gtest_main.cpp

@@ -0,0 +1,75 @@
+#include <cuda_runtime_api.h>
+#include <gtest/gtest.h>
+
+#include <vector>
+
+namespace {
+// Adapted from
+// https://github.com/NVIDIA/cutlass/blob/main/test/unit/common/filter_architecture.cpp#L78
+
+// generate gtest filter string based on device compute capability
+std::string gen_gtest_filter() {
+  static const int kMaxComputeCapability = 10000;
+  // Text filters for each kernel based on supported compute capability
+  struct Filter {
+    // Test filter string
+    char const* filter;
+
+    // Minimum and Maximum compute capability for the test group
+    int min_compute_capability;
+    int max_compute_capability;
+  };
+  static const std::vector<Filter> filters = {
+      {"SM70*", 70, 75},
+      {"SM75*", 75, kMaxComputeCapability},
+      {"SM80*", 80, kMaxComputeCapability},
+      {"SM89*", 89, 89},
+      {"SM90*", 90, 90},
+      {"SM100*", 100, 100},
+      {"SM120*", 120, 120}};
+
+  int compute_capability = 0;
+  cudaError_t err;
+  int cudaDeviceId;
+  err = cudaGetDevice(&cudaDeviceId);
+  if (cudaSuccess != err) {
+    std::cerr << "*** Warn: Could not detect active GPU device ID" << " ["
+              << cudaGetErrorString(err) << "]" << std::endl;
+  } else {
+    cudaDeviceProp device_properties;
+    err = cudaGetDeviceProperties(&device_properties, cudaDeviceId);
+    if (cudaSuccess != err) {
+      std::cerr << "*** Warn: Could not get device properties for GPU "
+                << cudaDeviceId << " [" << cudaGetErrorString(err) << "]"
+                << std::endl;
+    } else {
+      compute_capability =
+          (device_properties.major * 10) + device_properties.minor;
+    }
+  }
+  // Set negative test filters
+  std::stringstream ss;
+  ss << "-";
+  int i = 0;
+  for (const auto& filter : filters) {
+    if (compute_capability >= filter.min_compute_capability &&
+        compute_capability <= filter.max_compute_capability) {
+      // Skip filter for supported tests
+      continue;
+    }
+    // add separator if not the first filter
+    ss << (i++ ? ":" : "") << filter.filter;
+  }
+  return ss.str();
+}
+}  // namespace
+
+int main(int argc, char** argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  // honor --gtest_filter from commandline
+  if (::testing::GTEST_FLAG(filter).empty() ||
+      ::testing::GTEST_FLAG(filter) == "*") {
+    ::testing::GTEST_FLAG(filter) = gen_gtest_filter();
+  }
+  return RUN_ALL_TESTS();
+}

src/engine/CMakeLists.txt

@@ -42,5 +42,5 @@ cc_test(
   DEPS
     :engine
     absl::time
-    GTest::gtest_main
+    :gtest_main
 )

src/kernels/attention/CMakeLists.txt

@@ -63,7 +63,7 @@ cc_test(
   DEPS
     :attention.template
     absl::random_random
-    GTest::gtest_main
+    :gtest_main
     torch
 )
 
@@ -76,7 +76,7 @@ cc_test(
   DEPS
     :attention.template
     absl::random_random
-    GTest::gtest_main
+    :gtest_main
     torch
 )
 
@@ -88,7 +88,7 @@ cc_test(
   DEPS
     :attention.template
     absl::random_random
-    GTest::gtest_main
+    :gtest_main
     torch
 )
 

src/kernels/attention/attn_combine_kernel_test.cu

@@ -190,7 +190,7 @@ TEST_P(AttnCombineKernelTest, Combine) {
 }
 
 INSTANTIATE_TEST_SUITE_P(
-    Combine,
+    SM80,
     AttnCombineKernelTest,
     ::testing::Combine(::testing::Values(torch::kHalf,
                                          torch::kBFloat16),  // q_dtype

src/kernels/attention/sm80_mha_pagedkv_test.cu

@@ -229,7 +229,7 @@ TEST_P(MHAKernelPagedKVTest, PageKV) {
 }
 
 INSTANTIATE_TEST_SUITE_P(
-    MHA,
+    SM80,
     MHAKernelPagedKVTest,
     ::testing::Combine(
         ::testing::Values(1, 2, 4),                          // batch_size

src/kernels/attention/sm80_mla_pagedkv_test.cu

@@ -213,7 +213,7 @@ TEST_P(MLAKernelPagedKVTest, PageKV) {
 }
 
 INSTANTIATE_TEST_SUITE_P(
-    MLA,
+    SM80,
     MLAKernelPagedKVTest,
     ::testing::Combine(::testing::Values(torch::kHalf,
                                          torch::kBFloat16),  // q_dtype

src/kernels/attention/sm80_mla_test.cu

@@ -126,7 +126,7 @@ TEST_P(MLAKernelTest, MLA) {
 }
 
 INSTANTIATE_TEST_SUITE_P(
-    MLA,
+    SM80,
     MLAKernelTest,
     ::testing::Combine(::testing::Values(torch::kHalf,
                                          torch::kBFloat16),  // q_dtype

src/kernels/gemm/CMakeLists.txt

@@ -23,6 +23,6 @@ cc_test(
   DEPS
     :gemm.kernels
     absl::random_random
-    GTest::gtest_main
+    :gtest_main
     torch
 )

src/kernels/gemm/sm80_grouped_gemm_test.cu

@@ -181,7 +181,7 @@ TEST_P(GroupedGemmKernelTest, GEMM) {
 }
 
 INSTANTIATE_TEST_SUITE_P(
-    GEMM,
+    SM80,
     GroupedGemmKernelTest,
     ::testing::Combine(::testing::Values(torch::kHalf,
                                          torch::kBFloat16),  // dtype

src/kernels/gemm/tile_scheduler_test.cu

@@ -57,7 +57,7 @@ TEST_P(TileSchedulerTest, StaticPersistent) {
 }
 
 INSTANTIATE_TEST_SUITE_P(
-    TileScheduler,
+    SM80,
     TileSchedulerTest,
     ::testing::Combine(::testing::Values(1, 2),     // cluster_m
                        ::testing::Values(1, 2),     // cluster_n

src/kernels/moe/CMakeLists.txt

@@ -28,5 +28,5 @@ cc_test(
   DEPS
     :moe.kernels
     absl::random_random
-    GTest::gtest_main
+    :gtest_main
 )

src/kernels/moe/align_block_kernel_test.cu

@@ -185,7 +185,7 @@ TEST_P(AlignBlockTest, AlignBlock) {
 }
 
 INSTANTIATE_TEST_SUITE_P(
-    Moe,
+    SM80,
     AlignBlockTest,
     ::testing::Combine(::testing::Values(torch::kFloat),  // dtype
                        ::testing::Values(2,

src/kernels/moe/grouped_topk_sigmoid_kernel_test.cu

@@ -127,7 +127,7 @@ TEST_P(GroupedTopkSigmoidTest, TopkSoftmax) {
 }
 
 INSTANTIATE_TEST_SUITE_P(
-    Moe,
+    SM80,
     GroupedTopkSigmoidTest,
     ::testing::Combine(::testing::Values(torch::kFloat),  // dtype
                        ::testing::Values(10),             // n_tokens

src/kernels/moe/permutation_kernel_test.cu

@@ -206,7 +206,7 @@ TEST_P(PermuteTest, Mask) {
 }
 
 INSTANTIATE_TEST_SUITE_P(
-    Moe,
+    SM80,
     PermuteTest,
     ::testing::Combine(::testing::Values(torch::kFloat,
                                          torch::kHalf,

src/kernels/moe/topk_softmax_kernel_test.cu

@@ -47,7 +47,7 @@ TEST_P(TopkSoftmaxTest, TopkSoftmax) {
 }
 
 INSTANTIATE_TEST_SUITE_P(
-    Moe,
+    SM80,
     TopkSoftmaxTest,
     ::testing::Combine(
         ::testing::Values(torch::kFloat),               // dtype

src/kernels/triton/example/CMakeLists.txt

@@ -23,7 +23,7 @@ cc_test(
     kernel_test.cpp
   DEPS 
     triton_example.kernel
-    GTest::gtest_main
+    :gtest_main
     torch
     glog::glog
 )

src/layers/CMakeLists.txt

@@ -76,7 +76,7 @@ cc_test(
     :layers
     :state_dict
     absl::random_random
-    GTest::gtest_main
+    :gtest_main
 )
 
 add_subdirectory(attention)

src/layers/attention/CMakeLists.txt

@@ -34,5 +34,5 @@ cc_test(
   DEPS
     :attention
     absl::random_random
-    GTest::gtest_main
+    :gtest_main
 )

src/layers/moe/CMakeLists.txt

@@ -28,5 +28,5 @@ cc_test(
   DEPS
     :token_dispatcher
     absl::random_random
-    GTest::gtest_main
+    :gtest_main
 )

src/memory/CMakeLists.txt

@@ -36,5 +36,5 @@ cc_test(
   DEPS
     :memory
     absl::random_random
-    GTest::gtest_main
+    :gtest_main
 )

src/model_loader/CMakeLists.txt

@@ -22,7 +22,7 @@ cc_test(
     state_dict_test.cpp
   DEPS
     :state_dict
-    GTest::gtest_main
+    :gtest_main
   DATA 
     data/test.pth 
     data/test.safetensors

src/model_parallel/CMakeLists.txt

@@ -33,6 +33,6 @@ cc_test(
     process_group_test.cpp
   DEPS
     :process_group
-    GTest::gtest_main
+    :gtest_main
 )
 

src/quantization/CMakeLists.txt

@@ -41,7 +41,7 @@ cc_test(
   DEPS
     :quantization
     :state_dict
-    GTest::gtest_main
+    :gtest_main
   DATA 
     data/gptq_small.safetensors
     data/gptq.safetensors

src/request/CMakeLists.txt

@@ -33,5 +33,5 @@ cc_test(
     sequence_test.cpp
   DEPS
     :request
-    GTest::gtest_main
+    :gtest_main
 )

src/sampling/CMakeLists.txt

@@ -26,5 +26,5 @@ cc_test(
     logits_processor_test.cpp
   DEPS
     :sampler
-    GTest::gtest_main
+    :gtest_main
 )

src/scheduler/CMakeLists.txt

@@ -29,5 +29,5 @@ cc_library(
 #   DEPS
 #     :scheduler
 #     absl::strings
-#     GTest::gtest_main
+#     :gtest_main
 # )

src/speculative/CMakeLists.txt

@@ -27,5 +27,5 @@ cc_test(
   DEPS
     :speculative
     absl::strings
-    GTest::gtest_main
+    :gtest_main
 )

src/tokenizer/CMakeLists.txt

@@ -32,7 +32,7 @@ cc_test(
     tiktoken_tokenizer_test.cpp
   DEPS
     :tokenizer
-    GTest::gtest_main
+    :gtest_main
   DATA 
     data/tokenizer.model
     data/test.tiktoken