From 50d6808080b95d29de3529da5827ac6439f8ec97 Mon Sep 17 00:00:00 2001
From: Jianyu Huang <jianyuhuang@meta.com>
Date: Tue, 14 May 2024 20:08:26 -0700
Subject: [PATCH] Add GenAI build only option (#2584)

Summary:
Add GenAI build only option to save build time (e.g., we don't need TBE for gen ai models)

```
with-proxy CUDA_HOME=/usr/local/cuda-12.1 python setup.py install -DTORCH_CUDA_ARCH_LIST="8.0;9.0;9.0a" -DFBGEMM_GENAI_ONLY=ON -j 60 2>&1 | tee out.log
```

Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/2584

Reviewed By: jspark1105

Differential Revision: D57261898

Pulled By: jianyuh

fbshipit-source-id: 8abb4a4f68a4f015fc52adbf48cf663714361fc4
---
 fbgemm_gpu/CMakeLists.txt                     |  5 ++++-
 .../experimental/example/CMakeLists.txt       | 20 +++++++++++++++++++
 fbgemm_gpu/experimental/gen_ai/CMakeLists.txt | 20 +++++++++++++++++++
 fbgemm_gpu/setup.py                           |  4 ++++
 4 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/fbgemm_gpu/CMakeLists.txt b/fbgemm_gpu/CMakeLists.txt
index 288abfbb1..7f3191469 100644
--- a/fbgemm_gpu/CMakeLists.txt
+++ b/fbgemm_gpu/CMakeLists.txt
@@ -24,6 +24,7 @@ set(CMAKE_VERBOSE_MAKEFILE on)
 
 option(FBGEMM_CPU_ONLY  "Build FBGEMM_GPU without GPU support" OFF)
 option(USE_ROCM         "Build FBGEMM_GPU for ROCm" OFF)
+option(FBGEMM_GENAI_ONLY  "Build FBGEMM_GPU with GEN AI only support" OFF)
 
 if((NOT FBGEMM_CPU_ONLY) AND
    ((EXISTS "/opt/rocm/") OR (EXISTS $ENV{ROCM_PATH})) AND
@@ -85,7 +86,9 @@ include(${CMAKEMODULES}/RocmSetup.cmake)
 # Build FBGEMM_GPU (Main) Module
 ################################################################################
 
-include(FbgemmGpu.cmake)
+if(NOT FBGEMM_GENAI_ONLY)
+  include(FbgemmGpu.cmake)
+endif()
 
 ################################################################################
 # Build Experimental Modules
diff --git a/fbgemm_gpu/experimental/example/CMakeLists.txt b/fbgemm_gpu/experimental/example/CMakeLists.txt
index e8bc39996..457ab86b1 100644
--- a/fbgemm_gpu/experimental/example/CMakeLists.txt
+++ b/fbgemm_gpu/experimental/example/CMakeLists.txt
@@ -10,6 +10,26 @@ include(${CMAKEMODULES}/Utilities.cmake)
 # Target Sources
 ################################################################################
 
+if(FBGEMM_GENAI_ONLY)
+  set(fbgemm_sources_include_directories
+    # FBGEMM
+    ${FBGEMM}/include
+    # FBGEMM_GPU
+    ${CMAKE_CURRENT_SOURCE_DIR}../..
+    ${CMAKE_CURRENT_SOURCE_DIR}../../include
+    ${CMAKE_CURRENT_SOURCE_DIR}../../../include
+    # Third-party
+    ${THIRDPARTY}/asmjit/src
+    ${THIRDPARTY}/cpuinfo/include
+    ${THIRDPARTY}/cutlass/include
+    ${THIRDPARTY}/cutlass/tools/util/include)
+
+  set(third_party_include_directories
+    ${THIRDPARTY}/asmjit/src
+    ${THIRDPARTY}/cpuinfo/include
+    ${THIRDPARTY}/cutlass/include)
+endif()
+
 set(experimental_example_cpp_source_files
     src/cutlass_sgemm_nn.cu
     src/example_ops.cpp)
diff --git a/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt b/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt
index 007571ede..a5051e04b 100644
--- a/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt
+++ b/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt
@@ -8,6 +8,26 @@
 # Target Sources
 ################################################################################
 
+if(FBGEMM_GENAI_ONLY)
+  set(fbgemm_sources_include_directories
+    # FBGEMM
+    ${FBGEMM}/include
+    # FBGEMM_GPU
+    ${CMAKE_CURRENT_SOURCE_DIR}/../..
+    ${CMAKE_CURRENT_SOURCE_DIR}/../../include
+    ${CMAKE_CURRENT_SOURCE_DIR}/../../../include
+    # Third-party
+    ${THIRDPARTY}/asmjit/src
+    ${THIRDPARTY}/cpuinfo/include
+    ${THIRDPARTY}/cutlass/include
+    ${THIRDPARTY}/cutlass/tools/util/include)
+
+  set(third_party_include_directories
+    ${THIRDPARTY}/asmjit/src
+    ${THIRDPARTY}/cpuinfo/include
+    ${THIRDPARTY}/cutlass/include)
+endif()
+
 set(attention_ops_sources
     src/attention/attention.cpp
     src/attention/gqa_attn_splitk.cu)
diff --git a/fbgemm_gpu/setup.py b/fbgemm_gpu/setup.py
index 792c6d1cc..5105b97bb 100644
--- a/fbgemm_gpu/setup.py
+++ b/fbgemm_gpu/setup.py
@@ -241,6 +241,10 @@ def _get_cxx11_abi():
             print("[SETUP.PY] Building the CPU-ONLY variant of FBGEMM_GPU ...")
             cmake_args.append("-DFBGEMM_CPU_ONLY=ON")
 
+        if self.args.package_variant == "genai":
+            print("[SETUP.PY] Building the GENAI-ONLY variant of FBGEMM_GPU ...")
+            cmake_args.append("-DFBGEMM_GENAI_ONLY=ON")
+
         if self.args.nvml_lib_path:
             cmake_args.append(f"-DNVML_LIB_PATH={self.args.nvml_lib_path}")