fix msvc build

InternLM · Nov 2, 2023 · 6de4a37 · 6de4a37
1 parent c8eedef
commit 6de4a37
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/src/turbomind/kernels/decoder_multihead_attention/test_decoder_multihead_attention.cu b/src/turbomind/kernels/decoder_multihead_attention/test_decoder_multihead_attention.cu
@@ -10,6 +10,7 @@
 
 #include <iomanip>
 #include <numeric>
+#include <random>
 
 using namespace turbomind;
 
@@ -106,15 +107,15 @@ int main(int argc, char* argv[])
 
     DecoderMultiHeadAttentionParams<half> params{};
 
-    constexpr int kHeadNum    = 32;
-    constexpr int kHeadDim    = 128;
-    constexpr int KvHeadNum   = 32;
-    constexpr int kBatchSize  = 32;
-    constexpr int kContextLen = 7306;
-    // constexpr int kContextLen  = 1024;
+    constexpr int kHeadNum   = 32;
+    constexpr int kHeadDim   = 128;
+    constexpr int KvHeadNum  = 32;
+    constexpr int kBatchSize = 1;
+    // constexpr int kContextLen = 7306;
+    constexpr int kContextLen  = 1024;
     constexpr int kSequenceLen = kContextLen + 1;
     constexpr int kBlockSz     = 128;
-    constexpr int kTestIter    = 1;
+    constexpr int kTestIter    = 10;
     constexpr int kMaxSplitK   = 1;
 
     RNG rng{};
@@ -256,7 +257,7 @@ int main(int argc, char* argv[])
 
     std::vector<thrust::universal_vector<half>> outputs;
 
-    for (int i = 0; i < std::max(kTestIter, 10); ++i) {
+    for (int i = 0; i < std::max(kTestIter, 1); ++i) {
         DispatchDecoderMultiheadAttention<half>(params);
         if (auto err = cudaGetLastError(); err != cudaSuccess) {
             std::cout << cudaGetErrorString(err) << "\n";

diff --git a/src/turbomind/kernels/decoder_multihead_attention/test_utils.h b/src/turbomind/kernels/decoder_multihead_attention/test_utils.h
@@ -3,6 +3,7 @@
 #pragma once
 
 #include "decoder_multihead_attention.h"
+#include "src/turbomind/macro.h"
 #include <cuda_fp16.h>
 #include <memory>