[CI/Build][AMD] Install triton_kernels via CMake (#50328)
Signed-off-by: Ryan Rock <ryan.rock@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: depthfirst-app[bot] <184448029+depthfirst-app[bot]@users.noreply.github.com> Co-authored-by: root <root@smci355-ccs-aus-m11-05.cs-aus.dcgpu>
Ryan Rock committed
Jul 30, 2026 at 22:05 UTC
6724051c5766fe5bff71ed4b5df5c9dd590a2329
3 files changed
+20
-30
cmake/external_projects/triton_kernels.cmake
+11
-6
@@ -1,7 +1,5 @@
1
# Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels
2
3
-set(DEFAULT_TRITON_KERNELS_TAG "v3.5.1")
4
-
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})
@@ -12,13 +10,20 @@ if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
10
)
11
12
else()
15
- set(TRITON_GIT "https://github.com/triton-lang/triton.git")
16
- message (STATUS "[triton_kernels] Fetch from ${TRITON_GIT}:${DEFAULT_TRITON_KERNELS_TAG}")
13
+ if (VLLM_TARGET_DEVICE STREQUAL "rocm")
14
+ set(TRITON_GIT "https://github.com/ROCm/triton.git")
15
+ # Pinned from release/internal/3.6.x
16
+ set(TRITON_KERNELS_TAG "0f380657dbf3ee86eb57558ff71df24f03b5d4e7")
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
20
- GIT_REPOSITORY https://github.com/triton-lang/triton.git
21
- GIT_TAG ${DEFAULT_TRITON_KERNELS_TAG}
25
+ GIT_REPOSITORY ${TRITON_GIT}
26
+ GIT_TAG ${TRITON_KERNELS_TAG}
27
GIT_PROGRESS TRUE
28
SOURCE_SUBDIR python/triton_kernels/triton_kernels
29
)
docker/Dockerfile.rocm_base
-2
@@ -121,8 +121,6 @@ RUN cd triton \
121
&& if [ ! -f setup.py ]; then cd python; fi \
122
&& python3 setup.py bdist_wheel --dist-dir=dist \
123
&& mkdir -p /app/install && cp dist/*.whl /app/install
124
-RUN if [ -d triton/python/triton_kernels ]; then pip install build && cd triton/python/triton_kernels \
125
- && python3 -m build --wheel && cp dist/*.whl /app/install; fi
124
125
126
###
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+9
-22
@@ -80,28 +80,15 @@ def _patch_make_bitmatrix_metadata() -> None:
80
import triton.language as tl
81
82
try:
83
- if current_platform.is_rocm():
84
- from triton_kernels.tensor_details import bitmatrix as _bm
85
- from triton_kernels.tensor_details.bitmatrix import (
86
- BitmatrixMetadata,
87
- _keyed_add,
88
- cdiv,
89
- )
90
- from triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501
91
- sum_bitmatrix_rows,
92
- )
93
- else:
94
- from vllm.third_party.triton_kernels.tensor_details import (
95
- bitmatrix as _bm,
96
- )
97
- from vllm.third_party.triton_kernels.tensor_details.bitmatrix import (
98
- BitmatrixMetadata,
99
- _keyed_add,
100
- cdiv,
101
- )
102
- from vllm.third_party.triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501
103
- sum_bitmatrix_rows,
104
- )
83
+ from triton_kernels.tensor_details import bitmatrix as _bm
84
+ from triton_kernels.tensor_details.bitmatrix import (
85
+ BitmatrixMetadata,
86
+ _keyed_add,
87
+ cdiv,
88
+ )
89
+ from triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501
90
+ sum_bitmatrix_rows,
91
+ )
92
except ImportError:
93
return
94