| 1 | # Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels |
| 2 | |
| 3 | # Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to |
| 4 | # be directly set to the triton_kernels python directory. |
| 5 | if (DEFINED ENV{TRITON_KERNELS_SRC_DIR}) |
| 6 | message(STATUS "[triton_kernels] Fetch from $ENV{TRITON_KERNELS_SRC_DIR}") |
| 7 | FetchContent_Declare( |
| 8 | triton_kernels |
| 9 | SOURCE_DIR $ENV{TRITON_KERNELS_SRC_DIR} |
| 10 | ) |
| 11 | |
| 12 | else() |
| 13 | if (VLLM_TARGET_DEVICE STREQUAL "rocm") |
| 14 | set(TRITON_GIT "https://github.com/ROCm/triton.git") |
| 15 | # Pinned from release/internal/3.8.x |
| 16 | set(TRITON_KERNELS_TAG "669b31acc1dd1b3fd93286afd5db67f65d9f7557") |
| 17 | else() |
| 18 | set(TRITON_GIT "https://github.com/triton-lang/triton.git") |
| 19 | set(TRITON_KERNELS_TAG "v3.5.1") |
| 20 | endif() |
| 21 | message (STATUS "[triton_kernels] Fetch from ${TRITON_GIT}:${TRITON_KERNELS_TAG}") |
| 22 | FetchContent_Declare( |
| 23 | triton_kernels |
| 24 | # TODO (varun) : Fetch just the triton_kernels directory from Triton |
| 25 | GIT_REPOSITORY ${TRITON_GIT} |
| 26 | GIT_TAG ${TRITON_KERNELS_TAG} |
| 27 | GIT_PROGRESS TRUE |
| 28 | SOURCE_SUBDIR python/triton_kernels/triton_kernels |
| 29 | ) |
| 30 | endif() |
| 31 | |
| 32 | # Fetch content |
| 33 | FetchContent_MakeAvailable(triton_kernels) |
| 34 | |
| 35 | if (NOT triton_kernels_SOURCE_DIR) |
| 36 | message (FATAL_ERROR "[triton_kernels] Cannot resolve triton_kernels_SOURCE_DIR") |
| 37 | endif() |
| 38 | |
| 39 | if (DEFINED ENV{TRITON_KERNELS_SRC_DIR}) |
| 40 | set(TRITON_KERNELS_PYTHON_DIR "${triton_kernels_SOURCE_DIR}/") |
| 41 | else() |
| 42 | set(TRITON_KERNELS_PYTHON_DIR "${triton_kernels_SOURCE_DIR}/python/triton_kernels/triton_kernels/") |
| 43 | endif() |
| 44 | |
| 45 | message (STATUS "[triton_kernels] triton_kernels is available at ${TRITON_KERNELS_PYTHON_DIR}") |
| 46 | |
| 47 | add_custom_target(triton_kernels) |
| 48 | |
| 49 | # Ensure the vllm/third_party directory exists before installation |
| 50 | install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/triton_kernels\")") |
| 51 | |
| 52 | ## Copy .py files to install directory. |
| 53 | install(DIRECTORY |
| 54 | ${TRITON_KERNELS_PYTHON_DIR} |
| 55 | DESTINATION |
| 56 | vllm/third_party/triton_kernels/ |
| 57 | COMPONENT triton_kernels |
| 58 | FILES_MATCHING PATTERN "*.py") |