add rocm support for fp8 benchmarks

pytorch · Aug 10, 2024 · 8c8b4d7 · 8c8b4d7
1 parent c1d093a
commit 8c8b4d7
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt b/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt
@@ -28,10 +28,19 @@ set(attention_ops_sources
     src/attention/attention.cpp
     src/attention/gqa_attn_splitk.cu)
 
-set(quantize_ops_sources
-    src/quantize/cutlass_extensions.cu
-    src/quantize/quantize.cu
-    src/quantize/quantize.cpp)
+if(USE_ROCM)
+	set(quantize_ops_sources
+		src/quantize/ck_extensions/fp8_tensorwise_gemm.hip
+		src/quantize/ck_extensions/fp8_rowwise_gemm.hip
+		src/quantize/ck_extensions/fp8_blockwise_gemm.hip
+		src/quantize/quantize.cu
+		src/quantize/quantize.cpp)
+else()
+	set(quantize_ops_sources
+		src/quantize/cutlass_extensions.cu
+		src/quantize/quantize.cu
+		src/quantize/quantize.cpp)
+endif()
 
 set(comm_ops_sources
     src/comm/car.cu