Break up cutlass_extensions.cu, pt 4

Summary: - Break up cutlass_extensions.cu, pt 4 Differential Revision: D61170170
pytorch · Aug 12, 2024 · 3d145e9 · 3d145e9
1 parent 9b22d07
commit 3d145e9
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 40 deletions.
diff --git a/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt b/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt
@@ -39,11 +39,11 @@ if(USE_ROCM)
     src/quantize/quantize.cpp)
 else()
   set(quantize_ops_sources
-    src/quantize/cutlass_extensions.cu
     src/quantize/cutlass_extensions/f8f8bf16.cu
     src/quantize/cutlass_extensions/f8f8bf16_blockwise.cu
     src/quantize/cutlass_extensions/f8f8bf16_cublas.cu
     src/quantize/cutlass_extensions/f8f8bf16_rowwise.cu
+    src/quantize/cutlass_extensions/f8f8bf16_tensorwise.cu
     src/quantize/cutlass_extensions/i8i8bf16.cu
     src/quantize/cutlass_extensions/f8i4bf16_rowwise.cu
     src/quantize/cutlass_extensions/i8i8bf16_dynamic.cu

diff --git a/fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_blockwise.cu b/fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_blockwise.cu
@@ -19,8 +19,8 @@
 #include <cutlass/epilogue/collective/collective_builder.hpp> // @manual
 // clang-format on
 
+#include "cutlass_extensions/include/fp8_blockwise_cutlass_helpers.h"
 #include "cutlass_extensions/include/kernel_mode.h"
-#include "fp8_blockwise_cutlass_helpers.h"
 
 namespace {
 

diff --git a/...gen_ai/src/quantize/cutlass_extensions.cu → ...cutlass_extensions/f8f8bf16_tensorwise.cu b/...gen_ai/src/quantize/cutlass_extensions.cu → ...cutlass_extensions/f8f8bf16_tensorwise.cu
@@ -7,38 +7,9 @@
  */
 
 #include <ATen/ATen.h>
-#include <ATen/DeviceGuard.h>
-#include <ATen/Dispatch.h>
 #include <ATen/cuda/CUDAContext.h>
-#include <ATen/cuda/Exceptions.h>
-#include <ATen/cuda/Atomic.cuh>
-#if !(                                                  \
-    defined(USE_ROCM) ||                                \
-    ((defined(CUDA_VERSION) && CUDA_VERSION < 11000) || \
-     (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800))))
-#include <cublasLt.h>
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <cuda/atomic>
-#elif (defined(USE_ROCM))
-#include <hip/hip_bf16.h>
-#include <hip/hip_fp16.h>
-#include <hipblaslt/hipblaslt.h>
-#endif
-#include <c10/core/ScalarType.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <cutlass/core_io.h>
-#include <cutlass/cutlass.h>
-#include <cutlass/gemm/device/gemm.h>
-#include <cutlass/half.h>
-#include <cutlass/numeric_types.h>
-#include <cutlass/trace.h>
-#include <cutlass/util/host_tensor.h>
-#include "cublas_utils.h"
-
-#if CUDART_VERSION >= 12000
-#include <cuda_fp8.h>
-#endif
+#include <cutlass/util/device_memory.h>
+#include <cutlass/util/packed_stride.hpp>
 
 // clang-format off
 // The fixed ordering of the headers is required for CUTLASS 3.2+
@@ -48,14 +19,7 @@
 #include <cutlass/epilogue/collective/collective_builder.hpp> // @manual
 // clang-format on
 
-#include <cute/atom/mma_atom.hpp>
-#include <cutlass/gemm/dispatch_policy.hpp>
-#include <cutlass/gemm/kernel/gemm_universal.hpp>
-#include <cutlass/util/packed_stride.hpp>
-
 #include "cutlass_extensions/include/kernel_mode.h"
-#include "cutlass_extensions/include/threadblock.h"
-#include "fp8_blockwise_cutlass_helpers.h"
 
 namespace fbgemm_gpu {
 

diff --git a/.../quantize/fp8_blockwise_cutlass_helpers.h → ...s/include/fp8_blockwise_cutlass_helpers.h b/.../quantize/fp8_blockwise_cutlass_helpers.h → ...s/include/fp8_blockwise_cutlass_helpers.h