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: #2584

Reviewed By: jspark1105

Differential Revision: D57261898

Pulled By: jianyuh

fbshipit-source-id: 8abb4a4f68a4f015fc52adbf48cf663714361fc4

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

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)

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)

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}")