[fbgemm_gpu] Update Nova jobs

.github/scripts/nova_postscript.bash

@@ -25,6 +25,11 @@ echo "[NOVA] Current working directory: $(pwd)"
 # Record time for each step
 start_time=$(date +%s)
 
+echo "################################################################################"
+echo "Environment Variables:"
+printenv
+echo "################################################################################"
+
 # Collect PyTorch environment information
 collect_pytorch_env_info "${BUILD_ENV_NAME}"
 end_time=$(date +%s)
@@ -42,8 +47,13 @@ echo "[NOVA] Time taken to install wheel: ${runtime} seconds"
 # Test with PyTest
 echo "[NOVA] Current working directory: $(pwd)"
 if [[ $CU_VERSION = cu* ]]; then
-  echo "[NOVA] Testing the CUDA variant of FBGEMM_GPU ..."
-  export fbgemm_variant="cuda"
+  if [[ ${BUILD_TARGET} == "genai" ]]; then
+    echo "[NOVA] Testing the GenAI variant of FBGEMM_GPU ..."
+    export fbgemm_variant="genai"
+  else
+    echo "[NOVA] Testing the CUDA variant of FBGEMM_GPU ..."
+    export fbgemm_variant="cuda"
+  fi
 
 elif [[ $CU_VERSION = rocm* ]]; then
   echo "[NOVA] Testing the ROCm variant of FBGEMM_GPU ..."

.github/scripts/nova_prescript.bash

@@ -21,6 +21,11 @@ BUILD_ENV_NAME=${CONDA_ENV}
 # Record time for each step
 start_time=$(date +%s)
 
+echo "################################################################################"
+echo "Environment Variables:"
+printenv
+echo "################################################################################"
+
 # Display System Info
 print_system_info
 end_time=$(date +%s)
@@ -94,8 +99,13 @@ if [[ $CU_VERSION = cu* ]]; then
   start_time=${end_time}
   echo "[NOVA] Time taken to find NVML_LIB_PATH: ${runtime} seconds"
 
-  echo "[NOVA] Building the CUDA variant of FBGEMM_GPU ..."
-  export fbgemm_variant="cuda"
+  if [[ ${BUILD_TARGET} == "genai" ]]; then
+    echo "[NOVA] Building the GenAI variant of FBGEMM_GPU ..."
+    export fbgemm_variant="genai"
+  else
+    echo "[NOVA] Building the CUDA variant of FBGEMM_GPU ..."
+    export fbgemm_variant="cuda"
+  fi
 
 elif [[ $CU_VERSION = rocm* ]]; then
   echo "[NOVA] Building the ROCm variant of FBGEMM_GPU ..."

.github/workflows/build_wheels_genai_linux_aarch64.yml

@@ -0,0 +1,64 @@
+name: Build FBGEMM GenAI Aarch64 Linux Wheels
+
+on:
+    pull_request:
+    push:
+      branches:
+        - nightly
+        - main
+      tags:
+        # Release candidate tag look like: v1.11.0-rc1
+        - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
+        - v[0-9]+.[0-9]+.[0-9]+
+    workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  id-token: write
+  contents: read
+
+jobs:
+  generate-matrix:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: linux-aarch64
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      with-cuda: disable
+
+  build:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs: generate-matrix
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - repository: pytorch/FBGEMM
+            smoke-test-script: ""
+            pre-script: ../.github/scripts/nova_prescript.bash
+            post-script: ../.github/scripts/nova_postscript.bash
+            env-var-script: .github/scripts/nova_dir.bash
+            package-name: fbgemm_gpu
+    name: ${{ matrix.repository }}
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    with:
+      repository: ${{ matrix.repository }}
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: ${{ matrix.pre-script }}
+      post-script: ${{ matrix.post-script }}
+      package-name: ${{ matrix.package-name }}
+      build-target: genai
+      env-var-script: ${{ matrix.env-var-script }}
+      smoke-test-script: ${{ matrix.smoke-test-script }}
+      trigger-event: ${{ github.event_name }}
+      architecture: aarch64
+      setup-miniconda: false
+      timeout: 210

.github/workflows/build_wheels_genai_linux_x86.yml

@@ -0,0 +1,54 @@
+name: Build FBGEMM GenAI x86 Linux Wheels
+
+on:
+  pull_request:
+  push:
+    branches:
+      - nightly
+      - main
+    tags:
+      # Release candidate tag look like: v1.11.0-rc1
+      - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
+      - v[0-9]+.[0-9]+.[0-9]+
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  id-token: write
+  contents: read
+
+jobs:
+  generate-matrix:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: linux
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      with-cuda: enable
+      with-rocm: disable
+      with-cpu: disable
+
+  build:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs: generate-matrix
+    name: pytorch/FBGEMM
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    with:
+      repository: pytorch/FBGEMM
+      ref: ""
+      pre-script: ../.github/scripts/nova_prescript.bash
+      post-script: ../.github/scripts/nova_postscript.bash
+      smoke-test-script: ""
+      env-var-script: .github/scripts/nova_dir.bash
+      package-name: fbgemm_gpu
+      build-target: genai
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      trigger-event: ${{ github.event_name }}
+      timeout: 240

.github/workflows/build_wheels_linux_aarch64.yml

@@ -1,13 +1,11 @@
-name: Build Aarch64 Linux Wheels
+name: Build FBGEMM_GPU Aarch64 Linux Wheels
 
 on:
     pull_request:
     push:
       branches:
         - nightly
         - main
-        # Release candidate branch look like: v1.11.0-release
-        - v[0-9]+.[0-9]+.[0-9]+-release+
       tags:
         # Release candidate tag look like: v1.11.0-rc1
         - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
@@ -32,6 +30,7 @@ jobs:
       test-infra-repository: pytorch/test-infra
       test-infra-ref: main
       with-cuda: disable
+
   build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: generate-matrix

.github/workflows/build_wheels_linux_x86.yml

@@ -1,13 +1,11 @@
-name: Build x86 Linux Wheels
+name: Build FBGEMM_GPU x86 Linux Wheels
 
 on:
   pull_request:
   push:
     branches:
       - nightly
       - main
-      # Release candidate branch look like: v1.11.0-rc1
-      - v[0-9]+.[0-9]+.[0-9]+-release+
     tags:
       # Release candidate tag look like: v1.11.0-rc1
       - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
@@ -34,6 +32,7 @@ jobs:
       with-cuda: enable
       with-rocm: enable
       with-cpu: enable
+
   build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: generate-matrix

fbgemm_gpu/setup.py

@@ -116,8 +116,12 @@ def package_name(self) -> str:
             sys.exit(0)
 
         elif self.nova_flag() == 0:
-            # The package name is the same for all build variants in Nova
-            pass
+            # In Nova, we are publishing genai packages separately from the main
+            # fbgemm_gpu package, so if the package variant is genai, we need to
+            # update the package name accordingly.  Otherwise, the package name
+            # is the same for all other build variants in Nova
+            if self.args.package_variant == "genai":
+                pkg_name = "fbgemm_gpu_genai"
 
         else:
             # If running outside of Nova workflow context, append the channel