Skip to content

Commit 1212bae

Browse files
txustxus
committed
Merge branch 'thunderkittens'
2 parents f7a2a83 + 2b570bd commit 1212bae

8 files changed

Lines changed: 219 additions & 168 deletions

File tree

‎CMakeLists.txt‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,12 @@ FetchContent_Declare(
8787
GIT_TAG 55d53aa
8888
)
8989

90+
FetchContent_Declare(
91+
thunderkittens
92+
GIT_REPOSITORY https://github.com/HazyResearch/ThunderKittens.git
93+
GIT_TAG 00ea8a2
94+
)
95+
9096
# On NixOS with Clang, clangd uses a different resource-dir than clang++
9197
# Set the correct resource-dir for all C++ compilation
9298
if(CMAKE_CXX_COMPILER_ID STREQUAL "Clang" AND DEFINED CLANG_RESOURCE_DIR AND NOT CLANG_RESOURCE_DIR STREQUAL "")
@@ -95,7 +101,7 @@ endif()
95101

96102

97103
# Make dependencies available
98-
FetchContent_MakeAvailable(fmtlib json safetensors_cpp tokenizers_cpp)
104+
FetchContent_MakeAvailable(fmtlib json safetensors_cpp tokenizers_cpp thunderkittens)
99105

100106
# compiled library code
101107
add_subdirectory(src/tensor)

‎Makefile‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,8 @@ prepare_benchmark:
9595

9696
.PHONY: benchmark_tensor
9797
benchmark: prepare_benchmark
98-
@cmake --build build --target bm_tensor
98+
@cmake --build build --target bm_tensor --target test_tensor_cuda
99+
@ctest --test-dir build -R "TensorCUDATest\.MatmulBf16$$" --output-on-failure
99100
./build/benchmarks/tensor/bm_tensor
100101

101102
.PHONY: benchmark_llama

‎benchmarks/tensor/CMakeLists.txt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
if(BACKEND_CUDA)
2-
add_executable(bm_tensor cpu/bm_ops.cpp cuda/bm_ops.cpp)
2+
#add_executable(bm_tensor cpu/bm_ops.cpp cuda/bm_ops.cpp)
3+
add_executable(bm_tensor cuda/bm_ops.cpp)
34
target_link_libraries(bm_tensor PRIVATE tensor_core tensor_cpu tensor_cuda benchmark::benchmark_main)
45
else()
56
add_executable(bm_tensor cpu/bm_ops.cpp)

‎benchmarks/tensor/cuda/bm_ops.cpp‎

Lines changed: 88 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,13 @@ static void BM_CUDA_AddBf16(State& state) {
3030
state.range(1) * bytes_per_element * 2);
3131
}
3232

33-
BENCHMARK(BM_CUDA_AddBf16)
34-
->Args({8192, 2048})
35-
->Args({16384, 2048})
36-
->Args({65536, 2048})
37-
->Args({262144, 2048})
38-
->Unit(kMillisecond)
39-
->UseRealTime();
33+
// BENCHMARK(BM_CUDA_AddBf16)
34+
// ->Args({8192, 2048})
35+
// ->Args({16384, 2048})
36+
// ->Args({65536, 2048})
37+
// ->Args({262144, 2048})
38+
// ->Unit(kMillisecond)
39+
// ->UseRealTime();
4040

4141
static void BM_CUDA_SumFp32LastDim(State& state) {
4242
Tensor<float, CUDA> tensor(
@@ -58,11 +58,11 @@ static void BM_CUDA_SumFp32LastDim(State& state) {
5858
state.range(1) * bytes_per_element);
5959
}
6060

61-
BENCHMARK(BM_CUDA_SumFp32LastDim)
62-
->Args({16384, 2048})
63-
->Args({65536, 2048})
64-
->Unit(kMillisecond)
65-
->UseRealTime();
61+
// BENCHMARK(BM_CUDA_SumFp32LastDim)
62+
// ->Args({16384, 2048})
63+
// ->Args({65536, 2048})
64+
// ->Unit(kMillisecond)
65+
// ->UseRealTime();
6666

6767
static void BM_CUDA_SumFp32FirstDim(State& state) {
6868
Tensor<float, CUDA> tensor(
@@ -84,11 +84,11 @@ static void BM_CUDA_SumFp32FirstDim(State& state) {
8484
state.range(1) * bytes_per_element);
8585
}
8686

87-
BENCHMARK(BM_CUDA_SumFp32FirstDim)
88-
->Args({16384, 2048})
89-
->Args({65536, 2048})
90-
->Unit(kMillisecond)
91-
->UseRealTime();
87+
// BENCHMARK(BM_CUDA_SumFp32FirstDim)
88+
// ->Args({16384, 2048})
89+
// ->Args({65536, 2048})
90+
// ->Unit(kMillisecond)
91+
// ->UseRealTime();
9292

9393
static void BM_CUDA_MaskedFillBf16(State& state) {
9494
Tensor<bfloat16, CUDA> tensor(
@@ -114,11 +114,11 @@ static void BM_CUDA_MaskedFillBf16(State& state) {
114114
state.range(1) * bytes_per_element);
115115
}
116116

117-
BENCHMARK(BM_CUDA_MaskedFillBf16)
118-
->Args({16384, 2048})
119-
->Args({65536, 2048})
120-
->Unit(kMillisecond)
121-
->UseRealTime();
117+
// BENCHMARK(BM_CUDA_MaskedFillBf16)
118+
// ->Args({16384, 2048})
119+
// ->Args({65536, 2048})
120+
// ->Unit(kMillisecond)
121+
// ->UseRealTime();
122122

123123
static void BM_CUDA_CatBf16FirstDim(State& state) {
124124
Tensor<bfloat16, CUDA> tensor_a(
@@ -142,11 +142,11 @@ static void BM_CUDA_CatBf16FirstDim(State& state) {
142142
state.range(1) * bytes_per_element * 2);
143143
}
144144

145-
BENCHMARK(BM_CUDA_CatBf16FirstDim)
146-
->Args({16384, 2048})
147-
->Args({65536, 2048})
148-
->Unit(kMillisecond)
149-
->UseRealTime();
145+
// BENCHMARK(BM_CUDA_CatBf16FirstDim)
146+
// ->Args({16384, 2048})
147+
// ->Args({65536, 2048})
148+
// ->Unit(kMillisecond)
149+
// ->UseRealTime();
150150

151151
static void BM_CUDA_CatBf16LastDim(State& state) {
152152
Tensor<bfloat16, CUDA> tensor_a(
@@ -170,8 +170,64 @@ static void BM_CUDA_CatBf16LastDim(State& state) {
170170
state.range(1) * bytes_per_element * 2);
171171
}
172172

173-
BENCHMARK(BM_CUDA_CatBf16LastDim)
174-
->Args({16384, 2048})
175-
->Args({65536, 2048})
176-
->Unit(kMillisecond)
177-
->UseRealTime();
173+
// BENCHMARK(BM_CUDA_CatBf16LastDim)
174+
// ->Args({16384, 2048})
175+
// ->Args({65536, 2048})
176+
// ->Unit(kMillisecond)
177+
// ->UseRealTime();
178+
179+
static void BM_CUDA_MatmulBf16(State& state) {
180+
auto m = static_cast<size_t>(state.range(0));
181+
auto k = static_cast<size_t>(state.range(1));
182+
auto n = static_cast<size_t>(state.range(2));
183+
184+
Tensor<bfloat16, CUDA> tensor_a({m, k});
185+
Tensor<bfloat16, CUDA> tensor_b({k, n});
186+
187+
tensor_a.fill_(bfloat16(4.0));
188+
tensor_b.fill_(bfloat16(3.0));
189+
190+
auto a_v = tensor_a.view();
191+
auto b_v = tensor_b.view();
192+
193+
for (auto _ : state) {
194+
DoNotOptimize(matmul(a_v, b_v));
195+
}
196+
197+
int64_t flops = 0;
198+
flops += state.iterations() * 2 * m * k * n;
199+
state.counters["FLOPs"] = Counter(flops, Counter::kIsRate);
200+
auto bytes_per_element = 2;
201+
state.SetBytesProcessed(static_cast<int64_t>(state.iterations()) * m * k * k * n *
202+
bytes_per_element);
203+
}
204+
205+
BENCHMARK(BM_CUDA_MatmulBf16)->Args({16384, 8192, 16384})->Unit(kMillisecond)->UseRealTime();
206+
207+
static void BM_CUDA_MatmulFp32(State& state) {
208+
auto m = static_cast<size_t>(state.range(0));
209+
auto k = static_cast<size_t>(state.range(1));
210+
auto n = static_cast<size_t>(state.range(2));
211+
212+
Tensor<float, CUDA> tensor_a({m, k});
213+
Tensor<float, CUDA> tensor_b({k, n});
214+
215+
tensor_a.fill_(float(4.0));
216+
tensor_b.fill_(float(3.0));
217+
218+
auto a_v = tensor_a.view();
219+
auto b_v = tensor_b.view();
220+
221+
for (auto _ : state) {
222+
DoNotOptimize(matmul(a_v, b_v));
223+
}
224+
225+
int64_t flops = 0;
226+
flops += state.iterations() * 2 * m * k * n;
227+
state.counters["FLOPs"] = Counter(flops, Counter::kIsRate);
228+
auto bytes_per_element = 2;
229+
state.SetBytesProcessed(static_cast<int64_t>(state.iterations()) * m * k * k * n *
230+
bytes_per_element);
231+
}
232+
233+
BENCHMARK(BM_CUDA_MatmulFp32)->Args({16384, 8192, 16384})->Unit(kMillisecond)->UseRealTime();

‎src/tensor/cuda/CMakeLists.txt‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,13 @@ add_library(tensor_cuda STATIC storage.cu loader.cu ops.cu kernels/fill.cu kerne
66

77
configure_cuda_target(tensor_cuda)
88

9+
target_compile_definitions(tensor_cuda PRIVATE KITTENS_BLACKWELL)
10+
911
target_link_libraries(tensor_cuda
1012
PUBLIC
1113
tensor_core
1214
CUDA::cudart
15+
CUDA::cuda_driver
1316
CUDA::cublas
1417
CUDA::nvToolsExt
1518
PRIVATE
@@ -21,6 +24,11 @@ target_include_directories(tensor_cuda
2124
${PROJECT_SOURCE_DIR}/include
2225
)
2326

27+
target_include_directories(tensor_cuda
28+
SYSTEM PRIVATE
29+
"${FETCHCONTENT_BASE_DIR}/thunderkittens-src/include"
30+
)
31+
2432
source_group(
2533
TREE "${PROJECT_SOURCE_DIR}/include"
2634
PREFIX "Header Files"

0 commit comments

Comments
 (0)