MslkDefault.cmake

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

################################################################################
# Target Sources
################################################################################

glob_files_nohip(mslk_cpp_source_files_cpu
  csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_*.cpp
  csrc/conv/*.cpp
  csrc/gemm/*.cpp
  csrc/moe/*.cpp
  csrc/quantize/*.cpp
  csrc/attention/ck/fmha/*.cpp)

glob_files_nohip(mslk_cpp_source_files_gpu
  csrc/moe/*.cu
  csrc/quantize/*.cu)

# Include FB-internal sources into the build
if(BUILD_FB_CODE
  AND (MSLK_BUILD_VARIANT STREQUAL BUILD_VARIANT_CUDA)
  AND (CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0))
  BLOCK_PRINT("[FBPKG] INCLUDING FB-internal sources ...")

  glob_files_nohip(fb_only_sources_cpu
      fb/csrc/*/*.cpp)

  glob_files_nohip(fb_only_sources_gpu
      fb/csrc/*/*.cu)

  if(MSLK_FBPKG_BUILD)
    BLOCK_PRINT("[FBPKG] MSLK_FBPKG_BUILD is set.")

  else()
    BLOCK_PRINT(
      "[FBPKG] MSLK_FBPKG_BUILD is NOT set,"
      "certain FB-internal sources will be excluded from the build.")

    # To allow a directory, add it to the regex group below, e.g.:
    #   set(FB_ONLY_ALLOW_REGEX "fb/csrc/(my_kernel_1|my_kernel_2)")
    set(FB_ONLY_ALLOW_REGEX "^$")

    list(FILTER fb_only_sources_cpu
      INCLUDE REGEX "${FB_ONLY_ALLOW_REGEX}")

    list(FILTER fb_only_sources_gpu
      INCLUDE REGEX "${FB_ONLY_ALLOW_REGEX}")
  endif()

  list(APPEND mslk_cpp_source_files_cpu ${fb_only_sources_cpu})
  list(APPEND mslk_cpp_source_files_gpu ${fb_only_sources_gpu})

else()
  BLOCK_PRINT("[FBPKG] Will NOT be including FB-internal sources ...")

endif()

# CUDA-specific sources
file(GLOB_RECURSE mslk_cpp_source_files_cuda
  csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_*.cu
  csrc/conv/cutlass/*.cu
  csrc/conv/cutlass/**/*.cu
  csrc/gemm/cublas/*.cu
  csrc/gemm/cublas/**/*.cu
  csrc/gemm/cutlass/*.cu
  csrc/gemm/cutlass/**/*.cu
  csrc/quantize/cutlass/*.cu
  csrc/quantize/cutlass/**/*.cu)

# HIP-specific sources
file(GLOB_RECURSE mslk_cpp_source_files_hip
  csrc/gemm/ck/*.hip
  csrc/gemm/ck/**/*.hip
  csrc/quantize/ck/*.hip
  csrc/quantize/ck/**/*.hip)

# HIP FMHA sources - built separately for only the latest GPU architecture
# to reduce build time and binary size
file(GLOB_RECURSE mslk_hip_fmha_sources
  csrc/attention/ck/fmha/hip_fmha/**/*.cpp)

################################################################################
# Build FMHA Static Library (gfx942 only, ROCm 7.x+)
################################################################################

set(mslk_fmha_deps)
if(NOT DEFINED ENV{MSLK_BUILD_HIP_FMHA} OR NOT "$ENV{MSLK_BUILD_HIP_FMHA}" STREQUAL "0")
  if(MSLK_BUILD_VARIANT STREQUAL BUILD_VARIANT_ROCM
      AND mslk_hip_fmha_sources
      AND HIP_VERSION VERSION_GREATER_EQUAL "7.0")
    # Build HIP flags for FMHA targeting only gfx942
    set(HIP_HCC_FLAGS_FMHA ${HIP_HCC_FLAGS})
    list(FILTER HIP_HCC_FLAGS_FMHA EXCLUDE REGEX "--offload-arch=")
    list(APPEND HIP_HCC_FLAGS_FMHA
      --offload-arch=gfx942
      -DFMHA_LIMIT_MAX_HEADDIM_TO_256=1)

    message("Building FMHA for arch: gfx942 (HIP_VERSION=${HIP_VERSION})")

    # Mark sources for HIPCC compilation
    set_source_files_properties(${mslk_hip_fmha_sources} PROPERTIES
      HIP_SOURCE_PROPERTY_FORMAT 1)

    hip_include_directories("${mslk_include_directories}")

    hip_add_library(mslk_hip_fmha STATIC
      ${mslk_hip_fmha_sources}
      ${MSLK_HIP_HCC_LIBRARIES}
      HIPCC_OPTIONS ${HIP_HCC_FLAGS_FMHA})

    target_include_directories(mslk_hip_fmha PUBLIC
      ${MSLK_HIP_INCLUDE}
      ${ROCRAND_INCLUDE}
      ${ROCM_SMI_INCLUDE}
      ${mslk_include_directories})

    target_include_directories(mslk_hip_fmha PRIVATE
      ${TORCH_INCLUDE_DIRS}
      ${NCCL_INCLUDE_DIRS})

    set_target_properties(mslk_hip_fmha PROPERTIES
      POSITION_INDEPENDENT_CODE ON)

    set(mslk_fmha_deps mslk_hip_fmha)
  endif()
else()
  message("Skipping HIP FMHA build (MSLK_BUILD_HIP_FMHA=0)")
endif()

################################################################################
# Build Shared Library
################################################################################

gpu_cpp_library(
  PREFIX
    mslk
  TYPE
    SHARED
  INCLUDE_DIRS
    ${mslk_include_directories}
    ${CMAKE_CURRENT_SOURCE_DIR}/csrc/attention/cuda/cutlass_blackwell_fmha
  CPU_SRCS
    ${mslk_cpp_source_files_cpu}
  GPU_SRCS
    ${mslk_cpp_source_files_gpu}
  CUDA_SPECIFIC_SRCS
    ${mslk_cpp_source_files_cuda}
  HIP_SPECIFIC_SRCS
    ${mslk_cpp_source_files_hip}
  DEPS
    ${mslk_fmha_deps}
)

################################################################################
# Install Shared Library and Python Files
################################################################################

add_to_package(
  DESTINATION mslk
  TARGETS mslk)

install(
  DIRECTORY bench
  DESTINATION mslk)

if(BUILD_FB_CODE)
  install(
    DIRECTORY fb/mslk/attention/flash_attn
    DESTINATION mslk/fb/mslk/attention)
  install(
    DIRECTORY fb/external/quack
    DESTINATION mslk/fb/external)
endif()