skip tests that depend on cuda

vectorch-ai · Jun 30, 2024 · 031f13b · 031f13b
1 parent 8070796
commit 031f13b
Show file tree

Hide file tree

Showing 7 changed files with 64 additions and 0 deletions.
diff --git a/src/layers/activation_test.cpp b/src/layers/activation_test.cpp
@@ -36,6 +36,10 @@ TEST_P(ActivationTest, Basic) {
   const auto& [device, dtype, activation, in_features, out_features] =
       GetParam();
 
+  if (device.is_cuda() && !torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   auto input = torch::rand({in_features, out_features},
                            torch::dtype(dtype).device(device));
 

diff --git a/src/layers/attention/attention_test.cpp b/src/layers/attention/attention_test.cpp
@@ -92,6 +92,9 @@ TEST_P(AttentionPrefillTest, Varlen) {
                head_dim,
                scale,
                alibi] = GetParam();
+  if (device.is_cuda() && !torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
 
   absl::BitGen gen;
 
@@ -180,6 +183,10 @@ TEST_P(AttentionDecodeTest, KVCache) {
                head_dim,
                scale,
                alibi] = GetParam();
+  if (device.is_cuda() && !torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   // make sure kv_max_seq_len >= q_max_seq_len
   if (kv_max_seq_len < q_max_seq_len) {
     GTEST_SKIP() << "kv_max_seq_len < q_max_seq_len";

diff --git a/src/layers/normalization_test.cpp b/src/layers/normalization_test.cpp
@@ -37,6 +37,10 @@ TEST(NormalizationTest, LayerNorm) {
 }
 
 TEST(NormalizationTest, LayerNormKernel) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   // TODO: test other device and dtype combinations
   const auto dtype = torch::kHalf;
   const auto device = torch::kCUDA;
@@ -98,6 +102,10 @@ TEST(NormalizationTest, RMSNorm) {
 }
 
 TEST(NormalizationTest, RMSNormKernel) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   // TODO: test other device and dtype combinations
   const auto dtype = torch::kHalf;
   const auto device = torch::kCUDA;
@@ -125,6 +133,10 @@ TEST(NormalizationTest, RMSNormKernel) {
 }
 
 TEST(NormalizationTest, RMSNormResidualKernel) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   const auto dtype = torch::kHalf;
   const auto device = torch::kCUDA;
   const auto options = torch::dtype(dtype).device(device);

diff --git a/src/layers/pos_embedding_test.cpp b/src/layers/pos_embedding_test.cpp
@@ -116,6 +116,10 @@ TEST_P(PosEmbeddingTest, Rotary) {
               theta,
               interleaved,
               max_position_embeddings] = GetParam();
+  if (device.is_cuda() && !torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   const auto options = torch::dtype(dtype).device(device);
 
   // prepare inputs
@@ -194,6 +198,10 @@ TEST_P(PosEmbeddingKernelTest, Rotary) {
               interleaved,
               max_position_embeddings] = GetParam();
 
+  if (device.is_cuda() && !torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   const auto options = torch::dtype(dtype).device(device);
   // prepare inputs
   torch::Tensor query = torch::rand({num_tokens, n_heads, head_dim}, options);

diff --git a/src/memory/kv_cache_test.cpp b/src/memory/kv_cache_test.cpp
@@ -14,6 +14,10 @@ TEST(KVCacheTest, Empty) {
 }
 
 TEST(KVCacheTest, Basic) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   const int num_kv_heads = 32;
   const int head_dim = 128;
   const int block_size = 8;
@@ -57,6 +61,10 @@ TEST(KVCacheTest, Basic) {
 }
 
 TEST(KVCacheTest, Random) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   const int64_t num_kv_heads = 12;
   const int64_t head_dim = 128;
   const int64_t block_size = 4;

diff --git a/src/quantization/qlinear_impl_test.cpp b/src/quantization/qlinear_impl_test.cpp
@@ -23,6 +23,10 @@ TEST(QlinearTest, Basic) {
 }
 
 TEST(QlinearTest, ColumnParallelQuantLinear) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   const int64_t in_features = 4096;
   const int64_t out_features = 4096;
   QuantArgs quant_args;
@@ -57,6 +61,10 @@ TEST(QlinearTest, ColumnParallelQuantLinear) {
 }
 
 TEST(QlinearTest, RowParallelQuantLinear) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   const int64_t in_features = 4096;
   const int64_t out_features = 4096;
   QuantArgs quant_args;

diff --git a/src/sampling/logits_processor_test.cpp b/src/sampling/logits_processor_test.cpp
@@ -52,6 +52,9 @@ TEST(LogitsProcessorTest, Temperature) {
 }
 
 TEST(LogitsProcessorTest, TemperatureKernel) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
   // Test TemperatureLogitsProcessor
   torch::ScalarType dtype(torch::kHalf);
   torch::Device device(torch::kCUDA);
@@ -121,6 +124,9 @@ TEST(LogitsProcessorTest, FrequencyPresencePenalty) {
 }
 
 TEST(LogitsProcessorTest, FrequencyPresencePenaltyKernel) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
   // Test FrequencyPresencePenaltyLogitsProcessor
   torch::ScalarType dtype(torch::kHalf);
   torch::Device device(torch::kCUDA);
@@ -209,6 +215,9 @@ TEST(LogitsProcessorTest, RepetitionPenalty) {
 }
 
 TEST(LogitsProcessorTest, RepetitionPenaltyKernel) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
   // Test RepetitionPenaltyLogitsProcessor
   torch::ScalarType dtype(torch::kHalf);
   torch::Device device(torch::kCUDA);
@@ -244,6 +253,10 @@ TEST(LogitsProcessorTest, RepetitionPenaltyKernel) {
 }
 
 TEST(LogitsProcessorTest, TopK) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   // Set the random seed
   torch::manual_seed(100);
   torch::ScalarType dtype(torch::kHalf);
@@ -284,6 +297,10 @@ TEST(LogitsProcessorTest, TopK) {
 }
 
 TEST(LogitsProcessorTest, TopP) {
+  if (!torch::cuda::is_available()) {
+    GTEST_SKIP() << "CUDA not available, skipping test";
+  }
+
   // Set the random seed
   torch::manual_seed(100);
   torch::ScalarType dtype(torch::kHalf);