Add FB python sources into genai CMakeLists.txt (pytorch#978)

Summary:
Pull Request resolved: facebookresearch/FBGEMM#978

X-link: pytorch#3886

-  Add FB python sources into genai CMakeLists.txt

Reviewed By: renganxu

Differential Revision: D71930483

fbshipit-source-id: 7dc732fe607e5d23bd0a5d16d1ec7f6e118bbe68

Author: q10

.github/scripts/fbgemm_gpu_build.bash

@@ -357,11 +357,20 @@ __configure_fbgemm_gpu_build () {
 
   # Set debugging options
   if [ "$fbgemm_release_channel" != "release" ] || [ "$BUILD_DEBUG" -eq 1 ]; then
+    echo "[BUILD] Enabling debug features in the build ..."
     build_args+=(
       --debug
     )
   fi
 
+  # Set FB-only options
+  if [ "$BUILD_INCLUDE_FB_ONLY" -eq 1 ]; then
+    echo "[BUILD] Enabling build of FB-only code ..."
+    build_args+=(
+      --use_fb_only
+    )
+  fi
+
   # shellcheck disable=SC2145
   echo "[BUILD] FBGEMM_GPU build arguments have been set:  ${build_args[@]}"
 }

.github/workflows/fbgemm_gpu_ci_genai_generic_infra.yml

@@ -49,6 +49,7 @@ jobs:
       BUILD_ENV: build_binary
       BUILD_VARIANT: genai
       BUILD_CUDA_VERSION: ${{ matrix.cuda-version }}
+      BUILD_INCLUDE_FB_ONLY: 1
     continue-on-error: true
     strategy:
       fail-fast: false

fbgemm_gpu/CMakeLists.txt

@@ -26,7 +26,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)
-option(USE_FB_ONLY       "Build FBGEMM_GPU FB only operators" OFF)
+option(USE_FB_ONLY       "Build FBGEMM_GPU FB-only operators" OFF)
 
 if((NOT FBGEMM_CPU_ONLY) AND
    ((EXISTS "/opt/rocm/") OR (EXISTS $ENV{ROCM_PATH})) AND

fbgemm_gpu/bench/tbe/split_table_batched_embeddings_benchmark.py

@@ -12,7 +12,7 @@
 import os
 import tempfile
 from contextlib import nullcontext
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, Optional
 
 import click
 import numpy as np

fbgemm_gpu/experimental/gen_ai/CMakeLists.txt

@@ -28,10 +28,12 @@ file(GLOB experimental_gen_ai_cpp_source_files_gpu
 
 # Set the source file for FB only CPP
 if(USE_FB_ONLY)
-  file(GLOB fb_only_ops_sources
-      fb/src/*/*.cu
-      fb/src/*/*.cpp)
-  list(APPEND experimental_gen_ai_cpp_source_files_gpu ${fb_only_ops_sources})
+  file(GLOB fb_only_cxx_sources
+      fb/src/moe/*.cu
+      fb/src/moe/*.cpp
+      fb/src/trt_llm/*.cu
+      fb/src/trt_llm/*.cpp)
+  list(APPEND experimental_gen_ai_cpp_source_files_gpu ${fb_only_cxx_sources})
 endif()
 
 # CUDA-specific sources
@@ -82,3 +84,9 @@ add_to_package(
   DESTINATION fbgemm_gpu/experimental/gen_ai
   TARGETS fbgemm_gpu_experimental_gen_ai
   FILES ${experimental_gen_ai_python_source_files})
+
+if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/fb/gen_ai)
+  install(
+    DIRECTORY fb/gen_ai
+    DESTINATION fbgemm_gpu/experimental)
+endif()

fbgemm_gpu/experimental/gen_ai/fb/gen_ai/moe/__init__.py

@@ -4,9 +4,15 @@
 
 import torch
 
-torch.ops.load_library(
-    "//deeplearning/fbgemm/fbgemm_gpu/experimental/gen_ai:index_shuffling_ops_gpu"
-)
+try:
+    # In OSS, load the TARGET directly
+    torch.ops.load_library(
+        "//deeplearning/fbgemm/fbgemm_gpu/experimental/gen_ai:index_shuffling_ops_gpu"
+    )
+except Exception:
+    # In OSS, import the shared libraries
+    # pyre-ignore [21]
+    import fbgemm_gpu.experimental.gen_ai  # noqa: F401
 
 index_shuffling = torch.ops.fbgemm.index_shuffling  # noqa F401
 from .gather_scatter import (  # noqa F401

fbgemm_gpu/experimental/gen_ai/fb/gen_ai/trt_llm/fused_car.py

@@ -6,9 +6,15 @@
 
 import torch
 
-torch.ops.load_library(
-    "//deeplearning/fbgemm/fbgemm_gpu/experimental/gen_ai:trt_llm_ops_gpu"
-)
+try:
+    # In OSS, load the TARGET directly
+    torch.ops.load_library(
+        "//deeplearning/fbgemm/fbgemm_gpu/experimental/gen_ai:trt_llm_ops_gpu"
+    )
+except Exception:
+    # In OSS, import the shared libraries
+    # pyre-ignore [21]
+    import fbgemm_gpu.experimental.gen_ai  # noqa: F401
 
 
 def fused_one_shot_allreduce_residual_rms_norm_init(

fbgemm_gpu/setup.py

@@ -79,7 +79,7 @@ def from_args(cls, argv: List[str]):
         parser.add_argument(
             "--use_fb_only",
             action="store_true",
-            help="Build FB only operators.",
+            help="Build FB-only operators.",
         )
         parser.add_argument(
             "--cxxprefix",
@@ -93,7 +93,7 @@ def from_args(cls, argv: List[str]):
         print(f"[SETUP.PY] Other arguments: {other_args}")
         return FbgemmGpuBuild(setup_py_args, other_args)
 
-    def nova_flag(self) -> Optional[bool]:
+    def nova_flag(self) -> Optional[int]:
         if "BUILD_FROM_NOVA" in os.environ:
             if str(os.getenv("BUILD_FROM_NOVA")) == "0":
                 return 0